from __future__ import annotations

from fivetran_connector_sdk.constants import UPSERT_TYPE

from fivetran.operations import FivetranOperationSink
from metadata import Column
from state import (
    INCREMENTAL_DDL_WATERMARK,
    STATE_INCREMENTAL,
    STATE_SNAPSHOT,
    STATE_SYNC_TASK_ID,
    STATE_VERSION,
    ConnectorState,
)
from tests.pytest_compat import pytest


def test_fivetran_operation_sink_flushes_raw_writer_before_checkpoint() -> None:
    op_module = _FakeOperations()
    sink = FivetranOperationSink(op_module)
    writer = sink.raw_csv_record_writer(
        table="test_table",
        operation_type=UPSERT_TYPE,
        columns=(Column("id", "BIGINT"),),
        null_marker=r"\N",
        binary_encoding="base64",
    )

    writer(["1"])
    assert op_module.operation_stream.batches == []

    sink.checkpoint(state=_state())

    assert [[record.data["id"].long for record in batch] for batch in op_module.operation_stream.batches] == [[1]]
    assert op_module.checkpoints == [
        {
            "version": 1,
            "sync_task_id": "task",
            "snapshot": {
                "snapshot_tso": "0",
                "truncate_done": False,
                "current_file": None,
                "next_row": 0,
            },
            "incremental": {"ddl_watermark": {}, "checkpoint_ts": "0"},
        }
    ]

    sink.checkpoint(state=_state())

    assert [[record.data["id"].long for record in batch] for batch in op_module.operation_stream.batches] == [[1]]


def test_fivetran_operation_sink_uses_one_record_stream_for_multiple_encoders() -> None:
    op_module = _FakeOperations()
    sink = FivetranOperationSink(op_module)
    first_writer = sink.raw_csv_record_writer(
        table="first_table",
        operation_type=UPSERT_TYPE,
        columns=(Column("id", "BIGINT"),),
        null_marker=r"\N",
        binary_encoding="base64",
    )
    second_writer = sink.raw_csv_record_writer(
        table="second_table",
        operation_type=UPSERT_TYPE,
        columns=(Column("name", "VARCHAR(20)"),),
        null_marker=r"\N",
        binary_encoding="base64",
    )

    first_writer(["1"])
    second_writer(["Alice"])
    sink.checkpoint(state=_state())

    assert [[record.table_name for record in batch] for batch in op_module.operation_stream.batches] == [
        ["first_table", "second_table"]
    ]
    assert op_module.operation_stream.records == []


def test_fivetran_operation_sink_standard_mode_rejects_raw_writer() -> None:
    op_module = _FakeOperations()
    sink = FivetranOperationSink(op_module, use_batched_stream=False)

    with pytest.raises(RuntimeError, match="batched_stream"):
        sink.raw_csv_record_writer(
            table="test_table",
            operation_type=UPSERT_TYPE,
            columns=(Column("id", "BIGINT"),),
            null_marker=r"\N",
            binary_encoding="base64",
        )
    sink.upsert(table="test_table", data={"id": 1})

    assert op_module.operation_stream.records == []
    assert op_module.operation_stream.batches == []
    assert op_module.upserts == [("test_table", {"id": 1})]


class _FakeOperationStream:
    def __init__(self) -> None:
        self.records = []
        self.batches = []

    def add_record(self, record) -> None:
        self.records.append(record)

    def add_records(self, records) -> None:
        self.batches.append(records)


class _FakeOperations:
    def __init__(self) -> None:
        self.operation_stream = _FakeOperationStream()
        self.checkpoints = []
        self.upserts = []

    def upsert(self, *, table, data) -> None:
        self.upserts.append((table, data.copy()))

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


def _state() -> ConnectorState:
    return ConnectorState(
        data={
            STATE_VERSION: 1,
            STATE_SYNC_TASK_ID: "task",
            STATE_SNAPSHOT: {},
            STATE_INCREMENTAL: {INCREMENTAL_DDL_WATERMARK: {}},
        }
    )