from __future__ import annotations

from datetime import date, datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

import polars as pl

from .datasets import (
    DATASET_SPECS,
    DatasetName,
    DatasetSpec,
    ReadMode,
    RunStatus,
    WriteMode,
    empty_frame,
)
from .db import DbSettings, TimescaleDatasetStore


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


class InMemoryDatasetStore:
    def __init__(self) -> None:
        self._frames: dict[DatasetName, pl.DataFrame] = {}

    def init(self) -> None:
        return None

    def health(self) -> dict[str, Any]:
        return {"enabled": False, "status": "memory"}

    def write(self, dataset: DatasetName, frame: pl.DataFrame) -> int:
        if frame.is_empty():
            return 0
        existing = self._frames.get(dataset)
        self._frames[dataset] = (
            frame.clone()
            if existing is None or existing.is_empty()
            else pl.concat([existing, frame], how="vertical_relaxed")
        )
        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]
        frame = self._frames.get(dataset)
        if frame is None or frame.is_empty():
            return empty_frame(dataset)
        filtered = frame.clone()
        time_column = _time_column(spec)
        if start is not None and time_column in filtered.columns:
            filtered = filtered.filter(pl.col(time_column) >= pl.lit(start))
        if end is not None and time_column in filtered.columns:
            filtered = filtered.filter(pl.col(time_column) <= pl.lit(end))
        for column, value in (filters or {}).items():
            if column not in filtered.columns:
                return empty_frame(dataset)
            filtered = filtered.filter(pl.col(column) == value)
        mode = read_mode or spec.default_read_mode
        if mode == ReadMode.LATEST and spec.primary_key and not filtered.is_empty():
            filtered = (
                filtered.sort([*spec.primary_key, "run_ts"])
                .unique(subset=list(spec.primary_key), keep="last", maintain_order=True)
                .sort([time_column, "run_ts"])
            )
        elif not filtered.is_empty() and "run_ts" in filtered.columns:
            filtered = filtered.sort([time_column, "run_ts"])
        return filtered

    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:
        frame = self.read(dataset, filters=filters, read_mode=ReadMode.ALL_RUNS)
        if frame.is_empty():
            return None
        target_column = "logical_ts" if "logical_ts" in frame.columns else _time_column(DATASET_SPECS[dataset])
        if target_column not in frame.columns:
            return None
        return frame[target_column].max()


class DatasetStore:
    def __init__(
        self,
        root: Path | str,
        *,
        db_settings: DbSettings | None = None,
        db_backend: TimescaleDatasetStore | None = None,
    ) -> None:
        self.root = Path(root)
        self.db_settings = db_settings or DbSettings.from_env()
        self.db = db_backend or (
            TimescaleDatasetStore(self.db_settings) if self.db_settings.enabled else None
        )
        self.backend = self.db or InMemoryDatasetStore()

    def validate_frame(self, dataset: DatasetName, frame: pl.DataFrame) -> pl.DataFrame:
        spec = DATASET_SPECS[dataset]
        missing = [column for column in spec.schema if column not in frame.columns]
        if missing:
            raise ValueError(f"Missing columns for {dataset.value}: {missing}")
        return frame.select(
            [
                pl.col(column).cast(dtype, strict=False).alias(column)
                for column, dtype in spec.schema.items()
            ]
        )

    def write(
        self,
        dataset: DatasetName,
        frame: pl.DataFrame,
        *,
        logical_ts: datetime | None = None,
        run_id: str | None = None,
        run_ts: datetime | None = None,
        status: RunStatus = RunStatus.SUCCESS,
        write_mode: WriteMode = WriteMode.IMMUTABLE,
        stats: dict[str, Any] | None = None,
        error_reason: str | None = None,
    ) -> list[Path]:
        del status, stats, error_reason
        if frame.is_empty() and not frame.columns:
            return []
        effective_run_ts = (run_ts or datetime.now(tz=timezone.utc)).astimezone(timezone.utc)
        effective_run_id = run_id or uuid4().hex
        prepared = self._prepare_frame(
            frame,
            logical_ts=logical_ts,
            run_id=effective_run_id,
            run_ts=effective_run_ts,
            write_mode=write_mode,
        )
        validated = self.validate_frame(dataset, prepared)
        if validated.is_empty():
            return []
        self.backend.write(dataset, validated)
        return [Path(dataset.value) / effective_run_id]

    def read(
        self,
        dataset: DatasetName,
        start: Any | None = None,
        end: Any | None = None,
        filters: dict[str, Any] | None = None,
        read_mode: ReadMode | None = None,
    ) -> pl.DataFrame:
        frame = self.backend.read(
            dataset,
            start=start,
            end=end,
            filters=filters,
            read_mode=read_mode,
        )
        return empty_frame(dataset) if frame.is_empty() and not frame.columns else self.validate_frame(dataset, frame)

    def read_latest(
        self,
        dataset: DatasetName,
        *,
        filters: dict[str, Any] | None = None,
        read_mode: ReadMode | None = None,
    ) -> pl.DataFrame:
        frame = self.backend.read_latest(dataset, filters=filters, read_mode=read_mode)
        return empty_frame(dataset) if frame.is_empty() and not frame.columns else self.validate_frame(dataset, frame)

    def latest_logical_ts(
        self,
        dataset: DatasetName,
        filters: dict[str, Any] | None = None,
    ) -> datetime | None:
        return self.backend.latest_logical_ts(dataset, filters=filters)

    def db_init(self) -> None:
        if not self.db:
            raise ValueError("Database is not enabled")
        self.db.init()

    def db_backfill(
        self,
        dataset: DatasetName,
        *,
        start: datetime | None = None,
        end: datetime | None = None,
        filters: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        del dataset, start, end, filters
        raise RuntimeError("Parquet backfill has been removed; write datasets directly to Postgres")

    def db_backfill_all(self) -> list[dict[str, Any]]:
        raise RuntimeError("Parquet backfill has been removed; write datasets directly to Postgres")

    def db_verify(
        self,
        dataset: DatasetName,
        *,
        start: datetime | None = None,
        end: datetime | None = None,
        filters: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        if not self.db:
            raise ValueError("Database is not enabled")
        return self.db.verify_with_frame(dataset, None, start=start, end=end, filters=filters)

    def db_health(self) -> dict[str, Any]:
        return self.backend.health()

    def compact(
        self,
        dataset: DatasetName,
        *,
        start_date: date | None = None,
        end_date: date | None = None,
        filters: dict[str, Any] | None = None,
    ) -> dict[str, int]:
        del start_date, end_date
        if self.db:
            return {"partitions": 0, "rows": 0, "deleted_files": 0}
        spec = DATASET_SPECS[dataset]
        frame = self.read(dataset, filters=filters, read_mode=ReadMode.ALL_RUNS)
        if frame.is_empty():
            return {"partitions": 0, "rows": 0, "deleted_files": 0}
        time_column = _time_column(spec)
        if spec.primary_key:
            frame = (
                frame.sort([*spec.primary_key, "run_ts"])
                .unique(subset=list(spec.primary_key), keep="last", maintain_order=True)
                .sort([time_column, "run_ts"])
            )
        casted = self.validate_frame(dataset, frame)
        memory_backend = self.backend
        if isinstance(memory_backend, InMemoryDatasetStore):
            memory_backend._frames[dataset] = casted
        return {"partitions": 1, "rows": casted.height, "deleted_files": 0}

    def _prepare_frame(
        self,
        frame: pl.DataFrame,
        *,
        logical_ts: datetime | None,
        run_id: str,
        run_ts: datetime,
        write_mode: WriteMode,
    ) -> pl.DataFrame:
        prepared = frame
        if "ts" in prepared.columns and "date" not in prepared.columns:
            prepared = prepared.with_columns(pl.col("ts").dt.date().alias("date"))
        if "logical_ts" not in prepared.columns:
            if logical_ts is not None:
                prepared = prepared.with_columns(pl.lit(logical_ts).alias("logical_ts"))
            elif "ts" in prepared.columns:
                prepared = prepared.with_columns(pl.col("ts").alias("logical_ts"))
            else:
                raise ValueError("Frame must contain ts or an explicit logical_ts must be provided")
        if "run_id" not in prepared.columns:
            prepared = prepared.with_columns(pl.lit(run_id).alias("run_id"))
        if "run_ts" not in prepared.columns:
            prepared = prepared.with_columns(pl.lit(run_ts).alias("run_ts"))
        if "write_mode" not in prepared.columns:
            prepared = prepared.with_columns(pl.lit(write_mode.value).alias("write_mode"))
        return prepared
