from __future__ import annotations

from collections.abc import Callable, Sequence
from copy import deepcopy
from dataclasses import dataclass, field
from functools import wraps
from typing import Any, Protocol, cast

from fivetran.batched_stream import install_batched_operation_stream_support
from fivetran.sdk_records import (
    OperationRecordPlan,
    RawCSVRecordWriter,
    RawCSVRecordWriterFactory,
    RecordStream,
    build_raw_csv_record_writer,
    build_record_stream,
)
from metadata import Column
from observability import OperationStats, elapsed_ms, timer_start
from state import ConnectorState

# The SDK starts consuming Operations.operation_stream immediately after it
# starts the user update() thread. Installing this class-level extension during
# module import avoids a race where the consumer enters the original
# _build_next_batch() before the update configuration has been parsed.
install_batched_operation_stream_support()


def _checkpoint_payload(state: ConnectorState) -> dict[str, Any]:
    return state.to_dict()


class OperationSink(Protocol):
    def upsert(self, *, table: str, data: dict[str, Any]) -> None: ...

    def update(self, *, table: str, modified: dict[str, Any]) -> None: ...

    def delete(self, *, table: str, keys: dict[str, Any]) -> None: ...

    def truncate(self, *, table: str) -> None: ...

    def checkpoint(self, *, state: ConnectorState) -> None: ...


class _RawCSVRecordWriterProvider(Protocol):
    raw_csv_record_writer: RawCSVRecordWriterFactory


def raw_csv_record_writer_factory(sink: object) -> RawCSVRecordWriterFactory:
    return cast(_RawCSVRecordWriterProvider, sink).raw_csv_record_writer


class FivetranOperationSink:
    def __init__(self, op_module: Any, *, use_batched_stream: bool = True) -> None:
        self.op = op_module
        self._use_batched_stream = use_batched_stream
        self._record_stream: RecordStream | None = None

    def upsert(self, *, table: str, data: dict[str, Any]) -> None:
        self._flush_record_stream()
        self.op.upsert(table=table, data=data)

    def raw_csv_record_writer(
        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: Any | None = None,
    ) -> RawCSVRecordWriter:
        if not self._use_batched_stream:
            raise RuntimeError("raw CSV record writer requires batched_stream operation mode")

        writer = build_raw_csv_record_writer(
            plan=OperationRecordPlan(
                table=table,
                operation_type=operation_type,
                columns=columns,
                null_marker=null_marker,
                binary_encoding=binary_encoding,
                input_offset=input_offset,
                include_column_names=include_column_names,
                diagnostic_columns=diagnostic_columns,
            ),
            record_stream=self._shared_record_stream(),
        )
        return writer

    def update(self, *, table: str, modified: dict[str, Any]) -> None:
        self._flush_record_stream()
        self.op.update(table=table, modified=modified)

    def delete(self, *, table: str, keys: dict[str, Any]) -> None:
        self._flush_record_stream()
        self.op.delete(table=table, keys=keys)

    def truncate(self, *, table: str) -> None:
        self._flush_record_stream()
        self.op.truncate(table=table)

    def checkpoint(self, *, state: ConnectorState) -> None:
        self._flush_record_stream()
        self.op.checkpoint(state=_checkpoint_payload(state))

    def _shared_record_stream(self) -> RecordStream:
        stream = self._record_stream
        if stream is None:
            stream = build_record_stream(
                add_record=self.op.operation_stream.add_record,
                add_records=getattr(self.op.operation_stream, "add_records", None),
            )
            self._record_stream = stream
        return stream

    def _flush_record_stream(self) -> None:
        if self._record_stream is not None:
            self._record_stream.flush()


def timed_operation(operation: str) -> Callable[[Callable[..., None]], Callable[..., None]]:
    def decorator(method: Callable[..., None]) -> Callable[..., None]:
        @wraps(method)
        def wrapper(self: TimedOperationSink, *args: Any, **kwargs: Any) -> None:
            start = timer_start()
            try:
                method(self, *args, **kwargs)
            finally:
                self.stats.record(operation, elapsed_ms(start))

        return wrapper

    return decorator


class TimedOperationSink:
    def __init__(self, sink: OperationSink, stats: OperationStats | None = None) -> None:
        self.sink = sink
        self.stats = stats or OperationStats()

    @timed_operation("upsert")
    def upsert(self, *, table: str, data: dict[str, Any]) -> None:
        self.sink.upsert(table=table, data=data)

    def raw_csv_record_writer(
        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: Any | None = None,
    ) -> RawCSVRecordWriter:
        return raw_csv_record_writer_factory(self.sink)(
            table=table,
            operation_type=operation_type,
            columns=columns,
            null_marker=null_marker,
            binary_encoding=binary_encoding,
            input_offset=input_offset,
            include_column_names=include_column_names,
            diagnostic_columns=diagnostic_columns,
        )

    @timed_operation("update")
    def update(self, *, table: str, modified: dict[str, Any]) -> None:
        self.sink.update(table=table, modified=modified)

    @timed_operation("delete")
    def delete(self, *, table: str, keys: dict[str, Any]) -> None:
        self.sink.delete(table=table, keys=keys)

    @timed_operation("truncate")
    def truncate(self, *, table: str) -> None:
        self.sink.truncate(table=table)

    @timed_operation("checkpoint")
    def checkpoint(self, *, state: ConnectorState) -> None:
        self.sink.checkpoint(state=state)


class CheckpointTrackingOperationSink:
    def __init__(self, sink: OperationSink) -> None:
        self.sink = sink
        self.checkpoint_called = False

    def upsert(self, *, table: str, data: dict[str, Any]) -> None:
        self.sink.upsert(table=table, data=data)

    def raw_csv_record_writer(
        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: Any | None = None,
    ) -> RawCSVRecordWriter:
        return raw_csv_record_writer_factory(self.sink)(
            table=table,
            operation_type=operation_type,
            columns=columns,
            null_marker=null_marker,
            binary_encoding=binary_encoding,
            input_offset=input_offset,
            include_column_names=include_column_names,
            diagnostic_columns=diagnostic_columns,
        )

    def update(self, *, table: str, modified: dict[str, Any]) -> None:
        self.sink.update(table=table, modified=modified)

    def delete(self, *, table: str, keys: dict[str, Any]) -> None:
        self.sink.delete(table=table, keys=keys)

    def truncate(self, *, table: str) -> None:
        self.sink.truncate(table=table)

    def checkpoint(self, *, state: ConnectorState) -> None:
        self.checkpoint_called = True
        self.sink.checkpoint(state=state)


@dataclass
class RecordingOperationSink:
    calls: list[tuple[str, dict[str, Any]]] = field(default_factory=list)

    def upsert(self, *, table: str, data: dict[str, Any]) -> None:
        self.calls.append(("upsert", {"table": table, "data": data.copy()}))

    def update(self, *, table: str, modified: dict[str, Any]) -> None:
        self.calls.append(("update", {"table": table, "modified": modified.copy()}))

    def delete(self, *, table: str, keys: dict[str, Any]) -> None:
        self.calls.append(("delete", {"table": table, "keys": keys.copy()}))

    def truncate(self, *, table: str) -> None:
        self.calls.append(("truncate", {"table": table}))

    def checkpoint(self, *, state: ConnectorState) -> None:
        self.calls.append(("checkpoint", {"state": deepcopy(_checkpoint_payload(state))}))


__all__ = [
    "CheckpointTrackingOperationSink",
    "FivetranOperationSink",
    "OperationSink",
    "RecordingOperationSink",
    "TimedOperationSink",
    "raw_csv_record_writer_factory",
    "timed_operation",
]