from __future__ import annotations

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

from metadata import Column
from type_mapping import (
    BINARY_BASE_TYPES,
    DECIMAL_BASE_TYPES,
    FLOAT_BASE_TYPES,
    INTEGER_BASE_TYPES,
    TEXT_BASE_TYPES,
    base_source_type,
    build_value_coercer,
    coerce_binary,
    coerce_timestamp,
    coerce_value,
    is_source_type_syncable,
    is_unsigned_source_type,
)

type ValueDecoder = Callable[[str | None], Any]
type ColumnSlot = tuple[int, str]
type TimestampColumnSlot = tuple[int, str, bool]
type ColumnDecoder = tuple[int, str, ValueDecoder]
type DecodeSegment = tuple[int, Any]
type CompiledRowDecoder = Callable[[Sequence[str | None], dict[str, Any]], None]


class RowDecoder(Protocol):
    def __call__(self, raw_values: Sequence[str | None]) -> dict[str, Any]: ...


class RowDecoderInto(Protocol):
    def __call__(self, raw_values: Sequence[str | None], decoded: dict[str, Any]) -> None: ...


_RAW_TEXT_BASE_TYPES = TEXT_BASE_TYPES - {"SET"}
_RAW_EMPTY_NULL_BASE_TYPES = {"DATE", "DATETIME", "TIME", "VECTOR", "XML"}
_TRUE_VALUES = {"1", "true", "t", "yes", "y"}
_OP_RAW = 1
_OP_INT = 2
_OP_FLOAT = 3
_OP_BOOL = 4
_OP_ONE_BIT_BOOL = 5
_OP_TIMESTAMP = 6
_OP_JSON = 7
_OP_SET = 8
_OP_BINARY = 9
_OP_CONVERT = 10


def decode_named_row(
    raw_values: Mapping[str, str | None],
    columns: Iterable[Column],
    *,
    null_marker: str,
    binary_encoding: str,
) -> dict[str, Any]:
    return _decode_column_values(
        ((column, raw_values.get(column.name)) for column in columns),
        null_marker=null_marker,
        binary_encoding=binary_encoding,
    )


def decode_ordered_row(
    raw_values: Sequence[str | None],
    columns: Sequence[Column],
    *,
    null_marker: str,
    binary_encoding: str,
) -> dict[str, Any]:
    return _decode_column_values(
        zip(columns, raw_values, strict=True),
        null_marker=null_marker,
        binary_encoding=binary_encoding,
    )


def build_ordered_row_decoder(
    columns: Sequence[Column],
    *,
    null_marker: str,
    binary_encoding: str,
) -> RowDecoder:
    return CSVRowDecoder.new(
        columns,
        null_marker=null_marker,
        binary_encoding=binary_encoding,
    ).decode


@dataclass(frozen=True)
class CSVRowDecoder:
    decode_row: CompiledRowDecoder
    decode_segments: tuple[DecodeSegment, ...]
    required_raw_columns: tuple[ColumnSlot, ...]
    required_int_columns: tuple[ColumnSlot, ...]
    required_float_columns: tuple[ColumnSlot, ...]
    required_bool_columns: tuple[ColumnSlot, ...]
    required_one_bit_bool_columns: tuple[ColumnSlot, ...]
    required_timestamp_columns: tuple[TimestampColumnSlot, ...]
    required_json_columns: tuple[ColumnSlot, ...]
    required_set_columns: tuple[ColumnSlot, ...]
    required_binary_columns: tuple[ColumnDecoder, ...]
    column_decoders: tuple[ColumnDecoder, ...]

    @classmethod
    def new(
        cls,
        columns: Sequence[Column],
        *,
        null_marker: str,
        binary_encoding: str,
        input_offset: int = 0,
    ) -> CSVRowDecoder:
        required_raw_columns: list[ColumnSlot] = []
        required_int_columns: list[ColumnSlot] = []
        required_float_columns: list[ColumnSlot] = []
        required_bool_columns: list[ColumnSlot] = []
        required_one_bit_bool_columns: list[ColumnSlot] = []
        required_timestamp_columns: list[TimestampColumnSlot] = []
        required_json_columns: list[ColumnSlot] = []
        required_set_columns: list[ColumnSlot] = []
        required_binary_columns: list[ColumnDecoder] = []
        column_decoders: list[ColumnDecoder] = []

        for index, column in enumerate(columns):
            if not is_source_type_syncable(column.source_type, precision=column.precision):
                continue
            row_index = index + input_offset
            base = base_source_type(column.source_type)
            if not column.nullable:
                if _is_direct_raw_column(column, base):
                    required_raw_columns.append((row_index, column.name))
                    continue
                if _is_direct_int_column(column, base):
                    required_int_columns.append((row_index, column.name))
                    continue
                if base in FLOAT_BASE_TYPES:
                    required_float_columns.append((row_index, column.name))
                    continue
                if base in {"BOOL", "BOOLEAN"}:
                    required_bool_columns.append((row_index, column.name))
                    continue
                if _is_one_bit_boolean_column(column, base):
                    required_one_bit_bool_columns.append((row_index, column.name))
                    continue
                if base == "TIMESTAMP":
                    required_timestamp_columns.append((row_index, column.name, column.primary_key))
                    continue
                if base == "JSON":
                    required_json_columns.append((row_index, column.name))
                    continue
                if base == "SET":
                    required_set_columns.append((row_index, column.name))
                    continue
                if base in BINARY_BASE_TYPES:
                    required_binary_columns.append((row_index, column.name, _binary_value_decoder(binary_encoding)))
                    continue
            column_decoders.append(
                (
                    row_index,
                    column.name,
                    _value_decoder(column, null_marker=null_marker, binary_encoding=binary_encoding),
                )
            )
        required_raw = tuple(required_raw_columns)
        required_int = tuple(required_int_columns)
        required_float = tuple(required_float_columns)
        required_bool = tuple(required_bool_columns)
        required_one_bit_bool = tuple(required_one_bit_bool_columns)
        required_timestamp = tuple(required_timestamp_columns)
        required_json = tuple(required_json_columns)
        required_set = tuple(required_set_columns)
        required_binary = tuple(required_binary_columns)
        converters = tuple(column_decoders)
        decode_segments = _decode_segments(
            required_raw=required_raw,
            required_int=required_int,
            required_float=required_float,
            required_bool=required_bool,
            required_one_bit_bool=required_one_bit_bool,
            required_timestamp=required_timestamp,
            required_json=required_json,
            required_set=required_set,
            required_binary=required_binary,
            converters=converters,
        )
        return cls(
            decode_row=_compile_decode_row(decode_segments),
            decode_segments=decode_segments,
            required_raw_columns=required_raw,
            required_int_columns=required_int,
            required_float_columns=required_float,
            required_bool_columns=required_bool,
            required_one_bit_bool_columns=required_one_bit_bool,
            required_timestamp_columns=required_timestamp,
            required_json_columns=required_json,
            required_set_columns=required_set,
            required_binary_columns=required_binary,
            column_decoders=converters,
        )

    def decode(self, raw_values: Sequence[str | None]) -> dict[str, Any]:
        decoded: dict[str, Any] = {}
        self.decode_into(raw_values, decoded)
        return decoded

    def decode_into(self, raw_values: Sequence[str | None], decoded: dict[str, Any]) -> None:
        self.decode_row(raw_values, decoded)


def build_ordered_row_decoder_into(
    columns: Sequence[Column],
    *,
    null_marker: str,
    binary_encoding: str,
) -> RowDecoderInto:
    return CSVRowDecoder.new(
        columns,
        null_marker=null_marker,
        binary_encoding=binary_encoding,
    ).decode_into


def _compile_decode_row(decode_segments: tuple[DecodeSegment, ...]) -> CompiledRowDecoder:
    namespace = _decode_row_namespace()
    source = _build_decode_row_source(decode_segments, namespace)
    return _compile_decode_row_source(source, namespace)


def _decode_row_namespace() -> dict[str, Any]:
    return {
        "__builtins__": {},
        "binary_decoders": [],
        "column_decoders": [],
        "coerce_timestamp": coerce_timestamp,
        "load_json": json.loads,
        "to_float": float,
        "to_int": int,
        "true_values": _TRUE_VALUES,
    }


def _build_decode_row_source(decode_segments: tuple[DecodeSegment, ...], namespace: dict[str, Any]) -> str:
    lines = ["def decode_row(values, result):"]

    for operation, columns in decode_segments:
        if operation == _OP_RAW:
            for index, name in cast(tuple[ColumnSlot, ...], columns):
                lines.append(f"    result[{name!r}] = values[{index}]")
        elif operation == _OP_INT:
            for index, name in cast(tuple[ColumnSlot, ...], columns):
                lines.append(f"    result[{name!r}] = to_int(values[{index}])")
        elif operation == _OP_FLOAT:
            for index, name in cast(tuple[ColumnSlot, ...], columns):
                lines.append(f"    result[{name!r}] = to_float(values[{index}])")
        elif operation == _OP_BOOL:
            for index, name in cast(tuple[ColumnSlot, ...], columns):
                lines.append(f"    result[{name!r}] = values[{index}].lower() in true_values")
        elif operation == _OP_ONE_BIT_BOOL:
            for index, name in cast(tuple[ColumnSlot, ...], columns):
                lines.append(f"    result[{name!r}] = values[{index}] == '1'")
        elif operation == _OP_TIMESTAMP:
            for index, name, primary_key in cast(tuple[TimestampColumnSlot, ...], columns):
                lines.append(f"    result[{name!r}] = coerce_timestamp(values[{index}], primary_key={primary_key!r})")
        elif operation == _OP_JSON:
            for index, name in cast(tuple[ColumnSlot, ...], columns):
                lines.append(f"    result[{name!r}] = load_json(values[{index}])")
        elif operation == _OP_SET:
            for index, name in cast(tuple[ColumnSlot, ...], columns):
                raw_name = f"raw_{index}"
                lines.append(f"    {raw_name} = values[{index}]")
                lines.append(f"    result[{name!r}] = [] if {raw_name} == '' else {raw_name}.split(',')")
        elif operation == _OP_BINARY:
            for index, name, decode_value in cast(tuple[ColumnDecoder, ...], columns):
                decoder_index = _append_decoder(namespace, "binary_decoders", decode_value)
                lines.append(f"    result[{name!r}] = binary_decoders[{decoder_index}](values[{index}])")
        else:
            for index, name, decode_value in cast(tuple[ColumnDecoder, ...], columns):
                decoder_index = _append_decoder(namespace, "column_decoders", decode_value)
                lines.append(f"    result[{name!r}] = column_decoders[{decoder_index}](values[{index}])")

    if len(lines) == 1:
        lines.append("    pass")

    return "\n".join(lines)


def _append_decoder(namespace: dict[str, Any], name: str, decode_value: ValueDecoder) -> int:
    decoders = cast(list[ValueDecoder], namespace[name])
    decoder_index = len(decoders)
    decoders.append(decode_value)
    return decoder_index


def _compile_decode_row_source(source: str, namespace: dict[str, Any]) -> CompiledRowDecoder:
    namespace["binary_decoders"] = tuple(cast(list[ValueDecoder], namespace["binary_decoders"]))
    namespace["column_decoders"] = tuple(cast(list[ValueDecoder], namespace["column_decoders"]))
    exec(compile(source, "<row_decoder>", "exec"), namespace)
    return cast(CompiledRowDecoder, namespace["decode_row"])


def _decode_segments(
    *,
    required_raw: tuple[ColumnSlot, ...],
    required_int: tuple[ColumnSlot, ...],
    required_float: tuple[ColumnSlot, ...],
    required_bool: tuple[ColumnSlot, ...],
    required_one_bit_bool: tuple[ColumnSlot, ...],
    required_timestamp: tuple[TimestampColumnSlot, ...],
    required_json: tuple[ColumnSlot, ...],
    required_set: tuple[ColumnSlot, ...],
    required_binary: tuple[ColumnDecoder, ...],
    converters: tuple[ColumnDecoder, ...],
) -> tuple[DecodeSegment, ...]:
    segments: list[DecodeSegment] = []
    if required_raw:
        segments.append((_OP_RAW, required_raw))
    if required_int:
        segments.append((_OP_INT, required_int))
    if required_float:
        segments.append((_OP_FLOAT, required_float))
    if required_bool:
        segments.append((_OP_BOOL, required_bool))
    if required_one_bit_bool:
        segments.append((_OP_ONE_BIT_BOOL, required_one_bit_bool))
    if required_timestamp:
        segments.append((_OP_TIMESTAMP, required_timestamp))
    if required_json:
        segments.append((_OP_JSON, required_json))
    if required_set:
        segments.append((_OP_SET, required_set))
    if required_binary:
        segments.append((_OP_BINARY, required_binary))
    if converters:
        segments.append((_OP_CONVERT, converters))
    return tuple(segments)


def _value_decoder(
    column: Column,
    *,
    null_marker: str,
    binary_encoding: str,
) -> ValueDecoder:
    base = base_source_type(column.source_type)
    if base in _RAW_TEXT_BASE_TYPES:
        return _text_value_decoder(null_marker)
    if base in DECIMAL_BASE_TYPES or (base == "BIGINT" and is_unsigned_source_type(column.source_type)):
        return _empty_null_raw_value_decoder(null_marker)
    if base in _RAW_EMPTY_NULL_BASE_TYPES:
        return _empty_null_raw_value_decoder(null_marker)

    return build_value_coercer(
        column.source_type,
        null_marker=null_marker,
        binary_encoding=binary_encoding,
        precision=column.precision,
        primary_key=column.primary_key,
    )


def _is_direct_raw_column(column: Column, base: str) -> bool:
    return (
        base in _RAW_TEXT_BASE_TYPES
        or base in _RAW_EMPTY_NULL_BASE_TYPES
        or base in DECIMAL_BASE_TYPES
        or (base == "BIGINT" and is_unsigned_source_type(column.source_type))
    )


def _is_direct_int_column(column: Column, base: str) -> bool:
    if base == "BIGINT" and is_unsigned_source_type(column.source_type):
        return False
    if base == "TINYINT" and column.precision == 1:
        return False
    return base in INTEGER_BASE_TYPES


def _is_one_bit_boolean_column(column: Column, base: str) -> bool:
    return (base == "BIT" or base == "TINYINT") and column.precision == 1


def _binary_value_decoder(binary_encoding: str) -> ValueDecoder:
    def decode_value(raw_value: str | None) -> bytes:
        return coerce_binary(cast(str, raw_value), binary_encoding=binary_encoding)

    return decode_value


def _text_value_decoder(null_marker: str) -> ValueDecoder:
    def decode_value(raw_value: str | None) -> str | None:
        if raw_value is None or raw_value == null_marker:
            return None
        return raw_value

    return decode_value


def _empty_null_raw_value_decoder(null_marker: str) -> ValueDecoder:
    def decode_value(raw_value: str | None) -> str | None:
        if raw_value is None or raw_value == null_marker or raw_value == "":
            return None
        return raw_value

    return decode_value


def _decode_column_values(
    column_values: Iterable[tuple[Column, str | None]],
    *,
    null_marker: str,
    binary_encoding: str,
) -> dict[str, Any]:
    decoded: dict[str, Any] = {}
    for column, raw_value in column_values:
        if not is_source_type_syncable(column.source_type, precision=column.precision):
            continue
        decoded[column.name] = coerce_value(
            raw_value,
            column.source_type,
            null_marker=null_marker,
            binary_encoding=binary_encoding,
            precision=column.precision,
            primary_key=column.primary_key,
        )
    return decoded


__all__ = [
    "CSVRowDecoder",
    "RowDecoder",
    "RowDecoderInto",
    "build_ordered_row_decoder",
    "build_ordered_row_decoder_into",
    "decode_named_row",
    "decode_ordered_row",
]