from __future__ import annotations

from fivetran_connector_sdk import operations as sdk_operations
from fivetran_connector_sdk.constants import DELETE_TYPE, TABLES_COLUMNS_TYPES, UPDATE_TYPE, UPSERT_TYPE
from fivetran_connector_sdk.protos import common_pb2

from fivetran.sdk_records import (
    DiagnosticColumnSource,
    OperationRecordPlan,
    RawValueRef,
    build_raw_csv_operation_writers,
    build_raw_csv_record_encoder,
    build_raw_csv_record_writer,
    build_record_stream,
)
from metadata import Column
from row_decoder import CSVRowDecoder
from type_mapping import map_source_type_to_fivetran


def test_raw_csv_record_encoder_matches_sdk_mapping_for_upsert_scalars() -> None:
    table = "test_lineitem_record_encoder"
    columns = (
        Column("order_key", "BIGINT"),
        Column("part_key", "INT"),
        Column("quantity", "DECIMAL(15,2)"),
        Column("discount", "DOUBLE"),
        Column("ship_date", "DATE"),
        Column("comment", "VARCHAR(44)"),
    )
    raw_values = ["1", "1552", "17.00", "0.04", "1998-09-02", "regular package"]

    sdk_record = _sdk_snapshot_record(
        table=table,
        columns=columns,
        raw_values=raw_values,
        source_schema="test",
        source_table="lineitem",
        include_diagnostic_columns=True,
    )
    encode_record = build_raw_csv_record_encoder(
        OperationRecordPlan(
            table=table,
            operation_type=UPSERT_TYPE,
            columns=columns,
            null_marker=r"\N",
            binary_encoding="base64",
            diagnostic_columns=DiagnosticColumnSource(
                commit_ts=0,
                source_schema="test",
                source_table="lineitem",
            ),
        )
    )

    assert encode_record(raw_values) == sdk_record


def test_raw_csv_record_encoder_matches_sdk_mapping_for_complex_scalars() -> None:
    table = "test_complex_record_encoder"
    columns = (
        Column("active", "BOOLEAN"),
        Column("flag", "BIT", precision=1),
        Column("created_at", "DATETIME"),
        Column("commit_ts", "TIMESTAMP"),
        Column("zero_pk_ts", "TIMESTAMP", primary_key=True),
        Column("blob", "VARBINARY"),
        Column("optional_note", "VARCHAR(20)"),
    )
    raw_values = [
        "true",
        "1",
        "2024-01-02 03:04:05.123456",
        "2024-01-02 03:04:05.123456",
        "0000-00-00 00:00:00",
        "aGVsbG8=",
        r"\N",
    ]

    sdk_record = _sdk_snapshot_record(
        table=table,
        columns=columns,
        raw_values=raw_values,
        source_schema="test",
        source_table="complex",
        include_diagnostic_columns=False,
    )
    encode_record = build_raw_csv_record_encoder(
        OperationRecordPlan(
            table=table,
            operation_type=UPSERT_TYPE,
            columns=columns,
            null_marker=r"\N",
            binary_encoding="base64",
        )
    )

    assert encode_record(raw_values) == sdk_record


def test_raw_csv_record_encoder_writes_raw_json_and_nullable_values() -> None:
    table = "test_json_set_record_encoder"
    columns = (
        Column("payload", "JSON", nullable=False),
        Column("tags", "SET", nullable=False),
        Column("optional_ship_date", "DATE"),
        Column("optional_amount", "DECIMAL(15,2)"),
        Column("optional_note", "TEXT"),
        Column("optional_time", "TIME"),
    )
    raw_values = ['{"b":2,"a":[1]}', "A,B", "", "", r"\N", ""]

    encode_record = build_raw_csv_record_encoder(
        OperationRecordPlan(
            table=table,
            operation_type=UPSERT_TYPE,
            columns=columns,
            null_marker=r"\N",
            binary_encoding="base64",
        )
    )

    record = encode_record(raw_values)

    assert record.data["payload"].json == raw_values[0]
    assert record.data["tags"].json == '["A", "B"]'
    assert record.data["optional_ship_date"].null is True
    assert record.data["optional_amount"].null is True
    assert record.data["optional_note"].null is True
    assert record.data["optional_time"].null is True


def test_raw_csv_record_encoder_supports_incremental_offset_and_dynamic_diagnostics() -> None:
    table = "test_incremental_record_encoder"
    columns = (
        Column("id", "BIGINT"),
        Column("name", "VARCHAR(20)"),
    )
    raw_values = ["U", "worker", "test", "123", "1", "Alice"]

    sdk_record = _sdk_incremental_record(
        table=table,
        columns=columns,
        raw_values=raw_values,
        operation_type=UPDATE_TYPE,
        input_offset=4,
        diagnostic=True,
    )
    encode_record = build_raw_csv_record_encoder(
        OperationRecordPlan(
            table=table,
            operation_type=UPDATE_TYPE,
            columns=columns,
            input_offset=4,
            null_marker=r"\N",
            binary_encoding="base64",
            diagnostic_columns=DiagnosticColumnSource(
                commit_ts=RawValueRef(3),
                source_schema=RawValueRef(2),
                source_table=RawValueRef(1),
            ),
        )
    )

    assert encode_record(raw_values) == sdk_record


def test_raw_csv_record_encoder_supports_delete_primary_key_record() -> None:
    table = "test_delete_record_encoder"
    columns = (
        Column("id", "BIGINT", primary_key=True),
        Column("name", "VARCHAR(20)"),
    )
    raw_values = ["D", "worker", "test", "123", "1", "Alice"]

    sdk_record = _sdk_incremental_record(
        table=table,
        columns=columns,
        raw_values=raw_values,
        operation_type=DELETE_TYPE,
        input_offset=4,
        diagnostic=False,
        include_column_names=frozenset({"id"}),
    )
    encode_record = build_raw_csv_record_encoder(
        OperationRecordPlan(
            table=table,
            operation_type=DELETE_TYPE,
            columns=columns,
            input_offset=4,
            include_column_names=frozenset({"id"}),
            null_marker=r"\N",
            binary_encoding="base64",
        )
    )

    assert encode_record(raw_values) == sdk_record


def test_raw_csv_record_writer_batches_records() -> None:
    single_records = []
    batches = []
    stream = build_record_stream(
        add_record=single_records.append,
        add_records=batches.append,
        max_batch_records=2,
        max_batch_bytes=1_000_000,
    )
    writer = build_raw_csv_record_writer(
        plan=OperationRecordPlan(
            table="test_batch_writer",
            operation_type=UPSERT_TYPE,
            columns=(Column("id", "BIGINT"),),
            null_marker=r"\N",
            binary_encoding="base64",
        ),
        record_stream=stream,
    )

    writer(["1"])
    assert batches == []

    writer(["2"])
    assert len(batches) == 1
    assert [record.data["id"].long for record in batches[0]] == [1, 2]

    writer(["3"])
    writer.flush()

    assert single_records == []
    assert len(batches) == 2
    assert [record.data["id"].long for record in batches[1]] == [3]


def test_raw_csv_record_writer_can_share_record_stream() -> None:
    batches = []
    stream = build_record_stream(
        add_record=lambda record: None,
        add_records=batches.append,
        max_batch_records=10,
        max_batch_bytes=1_000_000,
    )
    writer = build_raw_csv_record_writer(
        plan=OperationRecordPlan(
            table="first_table",
            operation_type=UPSERT_TYPE,
            columns=(Column("id", "BIGINT"),),
            null_marker=r"\N",
            binary_encoding="base64",
        ),
        record_stream=stream,
    )
    other_writer = build_raw_csv_record_writer(
        plan=OperationRecordPlan(
            table="second_table",
            operation_type=UPSERT_TYPE,
            columns=(Column("name", "VARCHAR(20)"),),
            null_marker=r"\N",
            binary_encoding="base64",
        ),
        record_stream=stream,
    )

    writer(["1"])
    other_writer(["Alice"])
    stream.flush()

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


def test_raw_csv_operation_writers_preserve_mixed_operation_order() -> None:
    records = []
    stream = build_record_stream(
        add_record=records.append,
        add_records=records.extend,
        max_batch_records=100,
        max_batch_bytes=1_000_000,
    )

    def writer_factory(**kwargs):
        return build_raw_csv_record_writer(
            plan=OperationRecordPlan(**kwargs),
            record_stream=stream,
        )

    writers = build_raw_csv_operation_writers(
        writer_factory=writer_factory,
        table="test_ordered_operations",
        columns=(Column("id", "BIGINT", primary_key=True), Column("name", "VARCHAR(20)")),
        null_marker=r"\N",
        binary_encoding="base64",
        primary_key=("id",),
        include_delete=True,
    )

    assert writers is not None
    writers.write_upsert(["1", "Alice"])
    writers.write_delete(["2", "Bob"])
    writers.write_upsert(["3", "Carol"])
    writers.flush()

    assert [record.type for record in records] == [UPSERT_TYPE, DELETE_TYPE, UPSERT_TYPE]
    assert [record.data["id"].long for record in records] == [1, 2, 3]
    assert "name" not in records[1].data


def _sdk_snapshot_record(
    *,
    table: str,
    columns: tuple[Column, ...],
    raw_values: list[str],
    source_schema: str,
    source_table: str,
    include_diagnostic_columns: bool,
):
    TABLES_COLUMNS_TYPES[table] = _sdk_column_types(columns, include_diagnostic_columns=include_diagnostic_columns)
    decoded: dict[str, object] = {}
    CSVRowDecoder.new(columns, null_marker=r"\N", binary_encoding="base64").decode_into(raw_values, decoded)
    if include_diagnostic_columns:
        decoded["_incremental_commit_ts"] = 0
        decoded["_incremental_source_schema"] = source_schema
        decoded["_incremental_source_table"] = source_table
    return sdk_operations.connector_sdk_pb2.Record(
        schema_name=None,
        table_name=table,
        type=UPSERT_TYPE,
        data=sdk_operations._map_data_to_columns(decoded, sdk_operations._get_columns(table)),
    )


def _sdk_incremental_record(
    *,
    table: str,
    columns: tuple[Column, ...],
    raw_values: list[str],
    operation_type: int,
    input_offset: int,
    diagnostic: bool,
    include_column_names: frozenset[str] | None = None,
):
    TABLES_COLUMNS_TYPES[table] = _sdk_column_types(columns, include_diagnostic_columns=diagnostic)
    decoded: dict[str, object] = {}
    CSVRowDecoder.new(columns, null_marker=r"\N", binary_encoding="base64", input_offset=input_offset).decode_into(
        raw_values, decoded
    )
    if include_column_names is not None:
        decoded = {name: decoded[name] for name in include_column_names}
    if diagnostic:
        decoded["_incremental_commit_ts"] = int(raw_values[3])
        decoded["_incremental_source_schema"] = raw_values[2]
        decoded["_incremental_source_table"] = raw_values[1]
    return sdk_operations.connector_sdk_pb2.Record(
        schema_name=None,
        table_name=table,
        type=operation_type,
        data=sdk_operations._map_data_to_columns(decoded, sdk_operations._get_columns(table)),
    )


def _sdk_column_types(columns: tuple[Column, ...], *, include_diagnostic_columns: bool) -> dict[str, int]:
    result: dict[str, int] = {}
    for column in columns:
        fivetran_type = map_source_type_to_fivetran(column.source_type, precision=column.precision)
        if fivetran_type is not None:
            result[column.name] = getattr(common_pb2.DataType, fivetran_type)
    if include_diagnostic_columns:
        result["_incremental_commit_ts"] = common_pb2.DataType.LONG
        result["_incremental_source_schema"] = common_pb2.DataType.STRING
        result["_incremental_source_table"] = common_pb2.DataType.STRING
    return result