from __future__ import annotations

import json
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import Any, Literal, Protocol, cast

from fivetran_connector_sdk import operation_stream as sdk_operation_stream
from fivetran_connector_sdk.constants import DELETE_TYPE, UPSERT_TYPE
from fivetran_connector_sdk.protos import connector_sdk_pb2 as sdk_pb2

from metadata import Column
from type_mapping import (
    TEXT_BASE_TYPES,
    base_source_type,
    coerce_binary,
    is_source_type_syncable,
    map_source_type_to_fivetran,
)

type RawCSVRow = Sequence[str | None]
type AddRecord = Callable[[Any], None]
type AddRecords = Callable[[list[Any]], None]
type RecordEncoder = Callable[[RawCSVRow], Any]
type DiagnosticValueKind = Literal["long", "string"]

_TRUE_VALUES = {"1", "true", "t", "yes", "y"}
_SECONDS_PER_DAY = 86_400
_EPOCH_ORDINAL_OFFSET = 719_468
_DATE_SECONDS_CACHE_LIMIT = 8192


class RawCSVRecordWriter(Protocol):
    def __call__(self, raw_values: RawCSVRow) -> None: ...

    def flush(self) -> None: ...


class RecordStream(Protocol):
    def write(self, record: Any) -> None: ...

    def flush(self) -> None: ...


class _DateSecondsCache:
    __slots__ = ("limit", "values")

    def __init__(self, *, limit: int = _DATE_SECONDS_CACHE_LIMIT) -> None:
        self.limit = limit
        self.values: dict[str, int] = {}

    def seconds(self, raw: str) -> int:
        key = raw if len(raw) == 10 else raw[:10]
        values = self.values
        cached = values.get(key)
        if cached is not None:
            return cached
        seconds = _parse_date_seconds(key)
        if len(values) >= self.limit:
            values.clear()
        values[key] = seconds
        return seconds


@dataclass(frozen=True)
class RawValueRef:
    index: int


@dataclass(frozen=True)
class DiagnosticColumnSource:
    commit_ts: int | RawValueRef
    source_schema: str | RawValueRef
    source_table: str | RawValueRef


class RawCSVRecordWriterFactory(Protocol):
    def __call__(
        self,
        *,
        table: str,
        operation_type: int,
        columns: Sequence[Column],
        null_marker: str,
        binary_encoding: str,
        input_offset: int = 0,
        include_column_names: frozenset[str] | None = None,
        diagnostic_columns: DiagnosticColumnSource | None = None,
    ) -> RawCSVRecordWriter: ...


class RawCSVOperationWriters:
    __slots__ = ("delete", "upsert")

    def __init__(self, *, upsert: RawCSVRecordWriter, delete: RawCSVRecordWriter | None = None) -> None:
        self.upsert = upsert
        self.delete = delete

    def write_upsert(self, raw_values: RawCSVRow) -> None:
        self.upsert(raw_values)

    def write_delete(self, raw_values: RawCSVRow) -> None:
        if self.delete is None:
            raise ValueError("delete writer is not configured")
        self.delete(raw_values)

    def flush(self) -> None:
        self.upsert.flush()
        if self.delete is not None:
            self.delete.flush()


@dataclass(frozen=True)
class ColumnWritePlan:
    index: int
    name: str
    base: str
    fivetran_type: str
    nullable: bool
    empty_is_null: bool
    primary_key: bool


@dataclass(frozen=True)
class OperationRecordPlan:
    table: str
    operation_type: int
    columns: Sequence[Column]
    null_marker: str
    binary_encoding: str
    input_offset: int = 0
    include_column_names: frozenset[str] | None = None
    diagnostic_columns: DiagnosticColumnSource | None = None


def build_raw_csv_record_writer(
    *,
    plan: OperationRecordPlan,
    record_stream: RecordStream,
) -> RawCSVRecordWriter:
    return _EncodingRawCSVRecordWriter(
        encode=build_raw_csv_record_encoder(plan),
        stream=record_stream,
    )


def build_record_stream(
    *,
    add_record: AddRecord,
    add_records: AddRecords | None = None,
    max_batch_records: int = sdk_operation_stream.MAX_RECORDS_IN_BATCH,
    max_batch_bytes: int = sdk_operation_stream.MAX_BATCH_SIZE_IN_BYTES,
) -> RecordStream:
    if add_records is None:
        return _ImmediateRecordStream(add_record=add_record)
    return _BufferedRecordStream(
        add_records=add_records,
        max_batch_records=max_batch_records,
        max_batch_bytes=max_batch_bytes,
    )


def build_raw_csv_operation_writers(
    *,
    writer_factory: RawCSVRecordWriterFactory,
    table: str,
    columns: Sequence[Column],
    null_marker: str,
    binary_encoding: str,
    input_offset: int = 0,
    primary_key: Sequence[str] = (),
    diagnostic_columns: DiagnosticColumnSource | None = None,
    include_delete: bool = False,
) -> RawCSVOperationWriters:
    upsert_writer = writer_factory(
        table=table,
        operation_type=UPSERT_TYPE,
        columns=columns,
        null_marker=null_marker,
        binary_encoding=binary_encoding,
        input_offset=input_offset,
        diagnostic_columns=diagnostic_columns,
    )
    delete_writer = None
    if include_delete:
        delete_writer = writer_factory(
            table=table,
            operation_type=DELETE_TYPE,
            columns=columns,
            null_marker=null_marker,
            binary_encoding=binary_encoding,
            input_offset=input_offset,
            include_column_names=frozenset(primary_key),
        )
    return RawCSVOperationWriters(upsert=upsert_writer, delete=delete_writer)


class _EncodingRawCSVRecordWriter:
    __slots__ = ("_encode", "_flush_stream", "_write_record")

    def __init__(self, *, encode: RecordEncoder, stream: RecordStream) -> None:
        self._encode = encode
        self._write_record = stream.write
        self._flush_stream = stream.flush

    def __call__(self, raw_values: RawCSVRow) -> None:
        self._write_record(self._encode(raw_values))

    def flush(self) -> None:
        self._flush_stream()


class _ImmediateRecordStream:
    __slots__ = ("_add_record",)

    def __init__(self, *, add_record: AddRecord) -> None:
        self._add_record = add_record

    def write(self, record: Any) -> None:
        self._add_record(record)

    def flush(self) -> None:
        return


class _BufferedRecordStream:
    __slots__ = (
        "_add_records",
        "_batch_bytes",
        "_max_batch_bytes",
        "_max_batch_records",
        "_records",
    )

    def __init__(
        self,
        *,
        add_records: AddRecords,
        max_batch_records: int,
        max_batch_bytes: int,
    ) -> None:
        self._add_records = add_records
        self._max_batch_records = max(1, max_batch_records)
        self._max_batch_bytes = max(1, max_batch_bytes)
        self._records: list[Any] = []
        self._batch_bytes = 0

    def write(self, record: Any) -> None:
        record_size = record.ByteSize()
        records = self._records
        batch_bytes = self._batch_bytes
        if records and (len(records) >= self._max_batch_records or batch_bytes + record_size > self._max_batch_bytes):
            self.flush()
            records = self._records
            batch_bytes = self._batch_bytes
        records.append(record)
        batch_bytes += record_size
        self._batch_bytes = batch_bytes
        if len(records) >= self._max_batch_records or batch_bytes >= self._max_batch_bytes:
            self.flush()

    def flush(self) -> None:
        if not self._records:
            return
        records = self._records
        self._records = []
        self._batch_bytes = 0
        self._add_records(records)


def build_raw_csv_record_encoder(plan: OperationRecordPlan) -> RecordEncoder:
    columns = tuple(
        _column_write_plans(plan.columns, input_offset=plan.input_offset, include_names=plan.include_column_names)
    )
    namespace = _record_encoder_namespace(plan)
    source = _build_record_encoder_source(
        columns,
        diagnostic_columns=plan.diagnostic_columns,
        operation_type=plan.operation_type,
    )
    return _compile_record_encoder_source(source, namespace)


def _column_write_plans(
    columns: Sequence[Column],
    *,
    input_offset: int,
    include_names: frozenset[str] | None,
) -> Sequence[ColumnWritePlan]:
    plans: list[ColumnWritePlan] = []
    for index, column in enumerate(columns):
        if include_names is not None and column.name not in include_names:
            continue
        if not is_source_type_syncable(column.source_type, precision=column.precision):
            continue
        fivetran_type = map_source_type_to_fivetran(column.source_type, precision=column.precision)
        if fivetran_type is None:
            continue
        base = base_source_type(column.source_type)
        plans.append(
            ColumnWritePlan(
                index=index + input_offset,
                name=column.name,
                base=base,
                fivetran_type=fivetran_type,
                nullable=column.nullable,
                empty_is_null=base not in TEXT_BASE_TYPES,
                primary_key=column.primary_key,
            )
        )
    return tuple(plans)


def _record_encoder_namespace(plan: OperationRecordPlan) -> dict[str, Any]:
    return {
        "__builtins__": {},
        "Record": sdk_pb2.Record,
        "binary_encoding": plan.binary_encoding,
        "coerce_binary": coerce_binary,
        "date_seconds": _DateSecondsCache().seconds,
        "fraction_nanos": _fraction_nanos,
        "is_zero_datetime": _is_zero_datetime,
        "json_dumps": json.dumps,
        "null_marker": plan.null_marker,
        "operation_type": plan.operation_type,
        "table_name": plan.table,
        "time_seconds": _time_seconds,
        "to_float": float,
        "to_int": int,
        "true_values": _TRUE_VALUES,
    }


def _build_record_encoder_source(
    plans: Sequence[ColumnWritePlan],
    *,
    diagnostic_columns: DiagnosticColumnSource | None,
    operation_type: int,
) -> str:
    lines = [
        "def encode_record(values):",
        "    record = Record()",
        "    record.table_name = table_name",
    ]
    if operation_type != 0:
        lines.append("    record.type = operation_type")
    lines += [
        "    data = record.data",
    ]
    for position, plan in enumerate(plans):
        if plan.nullable:
            raw_name = f"raw_{position}"
            lines.append(f"    {raw_name} = values[{plan.index}]")
            null_condition = f"{raw_name} == null_marker"
            if plan.empty_is_null:
                null_condition += f" or {raw_name} == ''"
            lines.append(f"    if {null_condition}:")
            lines.append(f"        data[{plan.name!r}].null = True")
            lines.append("    else:")
            _append_value_write_lines(lines, plan, raw_name, "        ", position)
        else:
            _append_value_write_lines(lines, plan, f"values[{plan.index}]", "    ", position)

    if diagnostic_columns is not None:
        _append_diagnostic_value(lines, "_incremental_commit_ts", diagnostic_columns.commit_ts, "long")
        _append_diagnostic_value(lines, "_incremental_source_schema", diagnostic_columns.source_schema, "string")
        _append_diagnostic_value(lines, "_incremental_source_table", diagnostic_columns.source_table, "string")
    lines.append("    return record")
    return "\n".join(lines)


def _append_diagnostic_value(
    lines: list[str],
    column_name: str,
    value: int | str | RawValueRef,
    kind: DiagnosticValueKind,
) -> None:
    target = f"data[{column_name!r}].{kind}"
    if isinstance(value, RawValueRef):
        raw_name = f"diagnostic_{column_name}"
        if kind == "long":
            lines.append(f"    {raw_name} = to_int(values[{value.index}])")
            lines.append(f"    {target} = {raw_name}")
            return
        lines.append(f"    {raw_name} = values[{value.index}]")
        lines.append(f"    {target} = {raw_name}")
        return
    lines.append(f"    {target} = {value!r}")


def _append_value_write_lines(
    lines: list[str],
    plan: ColumnWritePlan,
    raw_name: str,
    indent: str,
    position: int,
) -> None:
    target = f"data[{plan.name!r}]"
    if plan.fivetran_type == "BOOLEAN":
        if plan.base in {"BOOL", "BOOLEAN"}:
            lines.append(f"{indent}{target}.bool = {raw_name}.lower() in true_values")
            return
        lines.append(f"{indent}{target}.bool = {raw_name} == '1'")
        return
    if plan.fivetran_type == "SHORT":
        lines.append(f"{indent}{target}.short = to_int({raw_name})")
        return
    if plan.fivetran_type == "INT":
        lines.append(f"{indent}{target}.int = to_int({raw_name})")
        return
    if plan.fivetran_type == "LONG":
        lines.append(f"{indent}{target}.long = to_int({raw_name})")
        return
    if plan.fivetran_type == "DECIMAL":
        lines.append(f"{indent}{target}.decimal = {raw_name}")
        return
    if plan.fivetran_type == "FLOAT":
        lines.append(f"{indent}{target}.float = to_float({raw_name})")
        return
    if plan.fivetran_type == "DOUBLE":
        lines.append(f"{indent}{target}.double = to_float({raw_name})")
        return
    if plan.fivetran_type == "NAIVE_DATE":
        lines.append(f"{indent}{target}.naive_date.seconds = date_seconds({raw_name})")
        return
    if plan.fivetran_type == "NAIVE_DATETIME":
        _append_datetime_lines(lines, f"{target}.naive_datetime", raw_name, indent, position)
        return
    if plan.fivetran_type == "UTC_DATETIME":
        _append_utc_datetime_lines(lines, target, raw_name, indent, position, primary_key=plan.primary_key)
        return
    if plan.fivetran_type == "BINARY":
        lines.append(f"{indent}{target}.binary = coerce_binary({raw_name}, binary_encoding=binary_encoding)")
        return
    if plan.fivetran_type == "XML":
        lines.append(f"{indent}{target}.xml = {raw_name}")
        return
    if plan.fivetran_type == "JSON":
        if plan.base == "SET":
            lines.append(f"{indent}{target}.json = json_dumps([] if {raw_name} == '' else {raw_name}.split(','))")
            return
        lines.append(f"{indent}{target}.json = {raw_name}")
        return
    if plan.fivetran_type == "STRING":
        lines.append(f"{indent}{target}.string = {raw_name}")
        return
    raise ValueError(f"unsupported Fivetran type for {plan.name}: {plan.fivetran_type}")


def _append_datetime_lines(
    lines: list[str],
    target: str,
    raw_name: str,
    indent: str,
    position: int,
) -> None:
    timestamp_name = f"timestamp_{position}"
    lines.append(f"{indent}{timestamp_name} = {target}")
    lines.append(f"{indent}{timestamp_name}.seconds = date_seconds({raw_name}) + time_seconds({raw_name})")
    lines.append(f"{indent}{timestamp_name}.nanos = fraction_nanos({raw_name})")


def _append_utc_datetime_lines(
    lines: list[str],
    target: str,
    raw_name: str,
    indent: str,
    position: int,
    *,
    primary_key: bool,
) -> None:
    lines.append(f"{indent}if is_zero_datetime({raw_name}):")
    if primary_key:
        lines.append(f"{indent}    {target}.utc_datetime.seconds = 0")
    else:
        lines.append(f"{indent}    {target}.null = True")
    lines.append(f"{indent}else:")
    _append_datetime_lines(lines, f"{target}.utc_datetime", raw_name, indent + "    ", position)


def _compile_record_encoder_source(source: str, namespace: dict[str, Any]) -> RecordEncoder:
    exec(compile(source, "<sdk_record_encoder>", "exec"), namespace)
    return cast(RecordEncoder, namespace["encode_record"])


def _is_zero_datetime(raw: str) -> bool:
    return raw.startswith("0000-00-00") and raw[11:19] == "00:00:00"


def _parse_date_seconds(raw: str) -> int:
    year = (ord(raw[0]) - 48) * 1000 + (ord(raw[1]) - 48) * 100 + (ord(raw[2]) - 48) * 10 + (ord(raw[3]) - 48)
    month = (ord(raw[5]) - 48) * 10 + (ord(raw[6]) - 48)
    day = (ord(raw[8]) - 48) * 10 + (ord(raw[9]) - 48)
    return _days_from_civil(year, month, day) * _SECONDS_PER_DAY


def _fraction_nanos(raw: str) -> int:
    if len(raw) <= 19 or raw[19] != ".":
        return 0
    digits = raw[20:]
    return int((digits + "000000000")[:9])


def _time_seconds(raw: str) -> int:
    hour = (ord(raw[11]) - 48) * 10 + (ord(raw[12]) - 48)
    minute = (ord(raw[14]) - 48) * 10 + (ord(raw[15]) - 48)
    second = (ord(raw[17]) - 48) * 10 + (ord(raw[18]) - 48)
    return hour * 3600 + minute * 60 + second


def _days_from_civil(year: int, month: int, day: int) -> int:
    year -= month <= 2
    era = (year if year >= 0 else year - 399) // 400
    year_of_era = year - era * 400
    day_of_year = (153 * (month + (-3 if month > 2 else 9)) + 2) // 5 + day - 1
    day_of_era = year_of_era * 365 + year_of_era // 4 - year_of_era // 100 + day_of_year
    return era * 146_097 + day_of_era - _EPOCH_ORDINAL_OFFSET


__all__ = [
    "DiagnosticColumnSource",
    "OperationRecordPlan",
    "RawCSVOperationWriters",
    "RawCSVRecordWriter",
    "RawCSVRecordWriterFactory",
    "RawValueRef",
    "RecordEncoder",
    "RecordStream",
    "build_raw_csv_operation_writers",
    "build_raw_csv_record_encoder",
    "build_raw_csv_record_writer",
    "build_record_stream",
]