from __future__ import annotations

import os
from dataclasses import dataclass
from datetime import date, datetime, timezone
from typing import Any, cast
from uuid import uuid4

import polars as pl
from psycopg import connect
from psycopg.rows import dict_row
from psycopg.types.json import Json

from .datasets import DATASET_SPECS, DatasetName, DatasetSpec, ReadMode, empty_frame

BACKFILL_START_SENTINEL = date(1970, 1, 1)
BACKFILL_END_SENTINEL = date(9999, 12, 31)


def _time_column(spec: DatasetSpec) -> str:
    return "ts" if "ts" in spec.schema else "logical_ts"


def _dtype_to_sql(dtype: Any) -> str:
    if dtype == pl.String:
        return "TEXT"
    if dtype == pl.Boolean:
        return "BOOLEAN"
    if dtype == pl.Int64:
        return "BIGINT"
    if dtype == pl.UInt64:
        return "NUMERIC(20,0)"
    if dtype == pl.Float64:
        return "DOUBLE PRECISION"
    if dtype == pl.Date:
        return "DATE"
    if dtype == pl.Datetime(time_zone="UTC"):
        return "TIMESTAMPTZ"
    return "TEXT"


def _normalize_value(value: Any) -> Any:
    if hasattr(value, "to_pydatetime"):
        value = value.to_pydatetime()
    if isinstance(value, datetime):
        if value.tzinfo is None:
            return value.replace(tzinfo=timezone.utc)
        return value.astimezone(timezone.utc)
    if isinstance(value, date):
        return value
    if isinstance(value, (dict, list)):
        return Json(value)
    return value


def _quoted(name: str) -> str:
    return '"' + name.replace('"', '""') + '"'


@dataclass(frozen=True)
class DbSettings:
    enabled: bool
    dsn: str | None
    schema: str
    read_mode: str
    write_mode: str
    strict_datasets: frozenset[str]

    @classmethod
    def from_env(cls) -> DbSettings:
        enabled = os.getenv("UFORIA_DB_ENABLED", "false").lower() == "true"
        dsn = os.getenv("UFORIA_DB_DSN")
        strict_raw = os.getenv("UFORIA_DB_STRICT_DATASETS", "")
        return cls(
            enabled=enabled and bool(dsn),
            dsn=dsn,
            schema=os.getenv("UFORIA_DB_SCHEMA", "uforia"),
            read_mode=os.getenv("UFORIA_DB_READ_MODE", "db"),
            write_mode=os.getenv("UFORIA_DB_WRITE_MODE", "db"),
            strict_datasets=frozenset(
                value.strip() for value in strict_raw.split(",") if value.strip()
            ),
        )


class TimescaleDatasetStore:
    def __init__(self, settings: DbSettings) -> None:
        self.settings = settings
        if not settings.dsn:
            raise ValueError("DB DSN is required when DB is enabled")
        self._column_cache: dict[DatasetName, frozenset[str]] = {}

    def _conn(self):
        return connect(self.settings.dsn, row_factory=dict_row)

    def table_name(self, dataset: DatasetName) -> str:
        return f"{self.settings.schema}.{dataset.value}"

    def _table_columns(self, dataset: DatasetName) -> frozenset[str]:
        cached = self._column_cache.get(dataset)
        if cached is not None:
            return cached
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT column_name
                    FROM information_schema.columns
                    WHERE table_schema = %s AND table_name = %s
                    """,
                    (self.settings.schema, dataset.value),
                )
                columns = frozenset(row["column_name"] for row in cur.fetchall())
        self._column_cache[dataset] = columns
        return columns

    def health(self) -> dict[str, Any]:
        try:
            with self._conn() as conn:
                with conn.cursor() as cur:
                    cur.execute("SELECT 1")
                    cur.fetchone()
            return {"enabled": True, "status": "ok"}
        except Exception as exc:  # pragma: no cover - runtime only
            return {"enabled": True, "status": "error", "detail": str(exc)}

    def init(self) -> None:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE SCHEMA IF NOT EXISTS {_quoted(self.settings.schema)}")
                timescaledb_available = False
                try:
                    cur.execute("CREATE EXTENSION IF NOT EXISTS timescaledb")
                    timescaledb_available = True
                except Exception:
                    conn.rollback()
                    cur = conn.cursor()
                    cur.execute(f"CREATE SCHEMA IF NOT EXISTS {_quoted(self.settings.schema)}")
                    cur.execute(
                        "SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'timescaledb')"
                    )
                    row = cur.fetchone()
                    timescaledb_available = bool(row["exists"]) if row else False
                for spec in DATASET_SPECS.values():
                    self._create_table(cur, spec, timescaledb_available=timescaledb_available)
                self._create_meta_tables(cur)
            conn.commit()

    def _create_meta_tables(self, cur: Any) -> None:
        schema = _quoted(self.settings.schema)
        cur.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {schema}._dataset_versions (
                dataset TEXT PRIMARY KEY,
                version BIGINT NOT NULL,
                updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
            )
            """
        )
        cur.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {schema}._backfill_state (
                dataset TEXT NOT NULL,
                start_date DATE,
                end_date DATE,
                asset TEXT NOT NULL DEFAULT '',
                venue TEXT NOT NULL DEFAULT '',
                status TEXT NOT NULL,
                rows_loaded BIGINT NOT NULL DEFAULT 0,
                updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                PRIMARY KEY (dataset, start_date, end_date, asset, venue)
            )
            """
        )
        for dataset, spec in DATASET_SPECS.items():
            cur.execute(
                f"""
                INSERT INTO {schema}._dataset_versions (dataset, version)
                VALUES (%s, %s)
                ON CONFLICT (dataset)
                DO UPDATE SET version = EXCLUDED.version, updated_at = NOW()
                """,
                (dataset.value, spec.version),
            )

    def _create_table(self, cur: Any, spec: DatasetSpec, *, timescaledb_available: bool) -> None:
        columns_sql = ", ".join(
            f"{_quoted(name)} {_dtype_to_sql(dtype)}" for name, dtype in spec.schema.items()
        )
        table = self.table_name(spec.name)
        cur.execute(f"CREATE TABLE IF NOT EXISTS {table} ({columns_sql})")
        cur.execute(
            """
            SELECT column_name
            FROM information_schema.columns
            WHERE table_schema = %s AND table_name = %s
            """,
            (self.settings.schema, spec.name.value),
        )
        existing_columns = {str(row["column_name"]) for row in cur.fetchall()}
        for name, dtype in spec.schema.items():
            if name in existing_columns:
                continue
            cur.execute(
                f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {_quoted(name)} {_dtype_to_sql(dtype)}"
            )
        self._column_cache.pop(spec.name, None)
        time_column = _time_column(spec)
        chunk_interval = "1 day" if spec.name in {
            DatasetName.MICROSTRUCTURE_RAW_EVENT,
            DatasetName.MICROSTRUCTURE_BOOK_STATE,
            DatasetName.MICROSTRUCTURE_BAR,
        } else "7 days"
        if timescaledb_available:
            cur.execute("SAVEPOINT hypertable_create")
            try:
                cur.execute(
                    """
                    SELECT create_hypertable(
                        %s,
                        %s,
                        chunk_time_interval => %s::interval,
                        if_not_exists => TRUE,
                        migrate_data => TRUE
                    )
                    """,
                    (table, time_column, chunk_interval),
                )
                cur.execute("RELEASE SAVEPOINT hypertable_create")
            except Exception:
                cur.execute("ROLLBACK TO SAVEPOINT hypertable_create")
                cur.execute("RELEASE SAVEPOINT hypertable_create")
            self._apply_timescale_policies(cur, spec, table)
        index_columns = list(
            dict.fromkeys([time_column, *spec.primary_key, *spec.partition_columns])
        )
        for idx, column in enumerate(index_columns):
            cur.execute(
                f"CREATE INDEX IF NOT EXISTS {_quoted(f'{spec.name.value}_{idx}_{column}_idx')} "
                f"ON {table} ({_quoted(column)})"
            )
        if spec.primary_key:
            unique_name = _quoted(f"{spec.name.value}_run_key")
            subset = ", ".join(_quoted(column) for column in [*spec.primary_key, "run_ts"])
            cur.execute(
                f"CREATE UNIQUE INDEX IF NOT EXISTS {unique_name} ON {table} ({subset})"
            )

    def _apply_timescale_policies(self, cur: Any, spec: DatasetSpec, table: str) -> None:
        raw_tables = {
            DatasetName.MICROSTRUCTURE_RAW_EVENT,
            DatasetName.MICROSTRUCTURE_BOOK_STATE,
            DatasetName.MICROSTRUCTURE_BAR,
        }
        if spec.name not in raw_tables:
            return
        cur.execute("SAVEPOINT timescale_policy")
        try:
            cur.execute(f"ALTER TABLE {table} SET (timescaledb.compress = true)")
            segment_cols = [
                column
                for column in ("asset", "venue", "event_type", "clock")
                if column in spec.schema
            ]
            if segment_cols:
                cur.execute(
                    f"ALTER TABLE {table} SET (timescaledb.compress_segmentby = %s)",
                    (",".join(segment_cols),),
                )
            cur.execute(
                (
                    "SELECT add_compression_policy("
                    "%s, compress_after => INTERVAL '7 days', if_not_exists => TRUE)"
                ),
                (table,),
            )
        except Exception:
            cur.execute("ROLLBACK TO SAVEPOINT timescale_policy")
        finally:
            cur.execute("RELEASE SAVEPOINT timescale_policy")

    def write(self, dataset: DatasetName, frame: pl.DataFrame) -> int:
        if frame.is_empty():
            return 0
        spec = DATASET_SPECS[dataset]
        columns = list(spec.schema.keys())
        insert_columns = ", ".join(_quoted(column) for column in columns)
        table = self.table_name(dataset)
        temp_table = _quoted(f"_tmp_{dataset.value}_{uuid4().hex}")
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE TEMP TABLE {temp_table} (LIKE {table}) ON COMMIT DROP")
                with cur.copy(f"COPY {temp_table} ({insert_columns}) FROM STDIN") as copy:
                    for row in frame.iter_rows(named=False):
                        copy.write_row(tuple(_normalize_value(value) for value in row))
                if spec.primary_key:
                    conflict_cols = ", ".join(
                        _quoted(column) for column in [*spec.primary_key, "run_ts"]
                    )
                    cur.execute(
                        f"""
                        INSERT INTO {table} ({insert_columns})
                        SELECT {insert_columns} FROM {temp_table}
                        ON CONFLICT ({conflict_cols}) DO NOTHING
                        """
                    )
                else:
                    cur.execute(
                        f"INSERT INTO {table} ({insert_columns}) "
                        f"SELECT {insert_columns} FROM {temp_table}"
                    )
            conn.commit()
        return frame.height

    def read(
        self,
        dataset: DatasetName,
        *,
        start: datetime | None = None,
        end: datetime | None = None,
        filters: dict[str, Any] | None = None,
        read_mode: ReadMode | None = None,
    ) -> pl.DataFrame:
        spec = DATASET_SPECS[dataset]
        filters = filters or {}
        mode = read_mode or spec.default_read_mode
        time_column = _time_column(spec)
        available_columns = self._table_columns(dataset)
        selected_names = [column for column in spec.schema if column in available_columns]
        if not selected_names:
            return empty_frame(dataset)
        select_columns = ", ".join(_quoted(column) for column in selected_names)
        params: list[Any] = []
        predicates: list[str] = []
        if start is not None:
            predicates.append(f"{_quoted(time_column)} >= %s")
            params.append(_normalize_value(start))
        if end is not None:
            predicates.append(f"{_quoted(time_column)} <= %s")
            params.append(_normalize_value(end))
        for column, value in filters.items():
            if column not in available_columns:
                return empty_frame(dataset)
            predicates.append(f"{_quoted(column)} = %s")
            params.append(_normalize_value(value))
        where_clause = f"WHERE {' AND '.join(predicates)}" if predicates else ""
        table = self.table_name(dataset)
        if mode == ReadMode.LATEST and spec.primary_key:
            latest_pk = [column for column in spec.primary_key if column in available_columns]
            if not latest_pk:
                return empty_frame(dataset)
            partition = ", ".join(_quoted(column) for column in latest_pk)
            sql = f"""
                SELECT {select_columns}
                FROM (
                    SELECT {select_columns},
                           row_number() OVER (
                               PARTITION BY {partition}
                               ORDER BY {_quoted('run_ts')} DESC
                           ) AS _rn
                    FROM {table}
                    {where_clause}
                ) latest_rows
                WHERE _rn = 1
                ORDER BY {_quoted(time_column)}, {_quoted('run_ts')}
            """
        else:
            sql = (
                f"SELECT {select_columns} FROM {table} {where_clause} "
                f"ORDER BY {_quoted(time_column)}, {_quoted('run_ts')}"
            )
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(sql, params)
                rows = cur.fetchall()
        if not rows:
            return empty_frame(dataset)
        frame = pl.from_dicts(rows, schema_overrides=spec.schema)
        for column, dtype in spec.schema.items():
            if column not in frame.columns:
                frame = frame.with_columns(pl.lit(None, dtype=dtype).alias(column))
        ordered = list(spec.schema)
        return frame.select(ordered)

    def read_latest(
        self,
        dataset: DatasetName,
        *,
        filters: dict[str, Any] | None = None,
        read_mode: ReadMode | None = None,
    ) -> pl.DataFrame:
        return self.read(dataset, filters=filters, read_mode=read_mode)

    def latest_logical_ts(
        self,
        dataset: DatasetName,
        filters: dict[str, Any] | None = None,
    ) -> datetime | None:
        spec = DATASET_SPECS[dataset]
        filters = filters or {}
        available_columns = self._table_columns(dataset)
        target_column = (
            "logical_ts"
            if "logical_ts" in available_columns
            else _time_column(spec) if _time_column(spec) in available_columns else None
        )
        if target_column is None:
            return None
        params: list[Any] = []
        predicates: list[str] = []
        for column, value in filters.items():
            if column not in available_columns:
                return None
            predicates.append(f"{_quoted(column)} = %s")
            params.append(_normalize_value(value))
        where_clause = f"WHERE {' AND '.join(predicates)}" if predicates else ""
        sql = (
            f"SELECT MAX({_quoted(target_column)}) AS latest_ts "
            f"FROM {self.table_name(dataset)} {where_clause}"
        )
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(sql, params)
                row = cur.fetchone()
        if not row:
            return None
        return cast(datetime | None, row["latest_ts"])

    def set_backfill_state(
        self,
        dataset: DatasetName,
        *,
        start_date: date | None,
        end_date: date | None,
        asset: str | None,
        venue: str | None,
        status: str,
        rows_loaded: int,
    ) -> None:
        normalized_start = start_date or BACKFILL_START_SENTINEL
        normalized_end = end_date or BACKFILL_END_SENTINEL
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"""
                    INSERT INTO {_quoted(self.settings.schema)}._backfill_state
                    (dataset, start_date, end_date, asset, venue, status, rows_loaded, updated_at)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, NOW())
                    ON CONFLICT (dataset, start_date, end_date, asset, venue)
                    DO UPDATE SET
                        status = EXCLUDED.status,
                        rows_loaded = EXCLUDED.rows_loaded,
                        updated_at = NOW()
                    """,
                    (
                        dataset.value,
                        normalized_start,
                        normalized_end,
                        asset or "",
                        venue or "",
                        status,
                        rows_loaded,
                    ),
                )
            conn.commit()

    def completed_backfill_dates(
        self,
        dataset: DatasetName,
        *,
        asset: str | None = None,
        venue: str | None = None,
    ) -> set[date]:
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    f"""
                    SELECT start_date
                    FROM {_quoted(self.settings.schema)}._backfill_state
                    WHERE dataset = %s
                      AND start_date = end_date
                      AND status = 'completed'
                      AND asset = %s
                      AND venue = %s
                    """,
                    (dataset.value, asset or "", venue or ""),
                )
                return {
                    row["start_date"]
                    for row in cur.fetchall()
                    if row["start_date"] not in {None, BACKFILL_START_SENTINEL}
                }

    def verify_with_frame(
        self,
        dataset: DatasetName,
        frame: pl.DataFrame | None = None,
        *,
        start: datetime | None = None,
        end: datetime | None = None,
        filters: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        spec = DATASET_SPECS[dataset]
        expected_frame = frame if frame is not None else empty_frame(dataset)
        db_frame = self.read(
            dataset,
            start=start,
            end=end,
            filters=filters,
            read_mode=ReadMode.ALL_RUNS,
        )
        db_latest = self.read(
            dataset,
            start=start,
            end=end,
            filters=filters,
            read_mode=ReadMode.LATEST,
        )
        latest_source = expected_frame
        if spec.primary_key and not expected_frame.is_empty():
            time_column = _time_column(spec)
            latest_source = (
                expected_frame.sort([*spec.primary_key, "run_ts"])
                .unique(subset=list(spec.primary_key), keep="last", maintain_order=True)
                .sort(time_column)
            )
        time_column = _time_column(DATASET_SPECS[dataset])
        source_min = (
            expected_frame[time_column].min()
            if time_column in expected_frame.columns and not expected_frame.is_empty()
            else None
        )
        source_max = (
            expected_frame[time_column].max()
            if time_column in expected_frame.columns and not expected_frame.is_empty()
            else None
        )
        db_min = (
            db_frame[time_column].min()
            if time_column in db_frame.columns and not db_frame.is_empty()
            else None
        )
        db_max = (
            db_frame[time_column].max()
            if time_column in db_frame.columns and not db_frame.is_empty()
            else None
        )
        return {
            "dataset": dataset.value,
            "source_all_runs_rows": expected_frame.height,
            "db_all_runs_rows": db_frame.height,
            "source_latest_rows": latest_source.height,
            "db_latest_rows": db_latest.height,
            "source_min_ts": source_min,
            "source_max_ts": source_max,
            "db_min_ts": db_min,
            "db_max_ts": db_max,
            "all_runs_row_delta": db_frame.height - expected_frame.height,
            "latest_row_delta": db_latest.height - latest_source.height,
        }
