from __future__ import annotations

import time
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any
from uuid import uuid4

import orjson
import polars as pl

from uforia.config import AppConfig, get_config
from uforia.ingest.deribit.client import DeribitPublicClient
from uforia.ingest.deribit.options import fetch_option_chain_payload, normalize_option_chain_payload
from uforia.storage import DatasetName, DatasetStore, ReadMode, RunStatus, WriteMode


@dataclass(frozen=True)
class IngestionAttemptResult:
    logical_ts: datetime
    run_id: str
    status: RunStatus
    rows_written_raw: int
    rows_written_normalized: int
    error_message: str | None = None


def floor_hour(ts: datetime) -> datetime:
    return ts.astimezone(timezone.utc).replace(minute=0, second=0, microsecond=0)


def current_collectible_hour(now: datetime, grace_seconds: int) -> datetime:
    boundary = floor_hour(now)
    if now.astimezone(timezone.utc) >= boundary + timedelta(seconds=grace_seconds):
        return boundary
    return boundary - timedelta(hours=1)


def archive_deribit_options_once(
    store: DatasetStore,
    *,
    as_of: datetime,
    force: bool = False,
    underlying: str = "BTC",
    venue: str = "deribit",
    client: DeribitPublicClient | None = None,
    config: AppConfig | None = None,
) -> IngestionAttemptResult:
    cfg = config or get_config()
    deribit = client or DeribitPublicClient(cfg)
    started_at = deribit.now().astimezone(timezone.utc)
    logical_ts = floor_hour(as_of)
    run_id = uuid4().hex

    if not force and _has_successful_snapshot(
        store, logical_ts, venue=venue, underlying=underlying
    ):
        _write_ingestion_run(
            store,
            logical_ts=logical_ts,
            run_id=run_id,
            run_ts=started_at,
            dataset=DatasetName.OPTION_CHAIN_SNAPSHOT.value,
            venue=venue,
            underlying=underlying,
            status=RunStatus.SKIPPED_EXISTING,
            rows_written_raw=0,
            rows_written_normalized=0,
            duration_ms=0,
        )
        return IngestionAttemptResult(
            logical_ts=logical_ts,
            run_id=run_id,
            status=RunStatus.SKIPPED_EXISTING,
            rows_written_raw=0,
            rows_written_normalized=0,
        )

    try:
        collectible = current_collectible_hour(started_at, cfg.ingestion.archive_grace_seconds)
        if logical_ts == collectible:
            payload = fetch_option_chain_payload(
                underlying=underlying,
                venue=venue,
                client=deribit,
                config=cfg,
            )
            raw_rows = _write_raw_snapshot(
                store,
                payload=payload,
                logical_ts=logical_ts,
                run_id=run_id,
                run_ts=started_at,
            )
        else:
            archived_payload = _load_archived_payload(
                store,
                logical_ts=logical_ts,
                venue=venue,
                underlying=underlying,
            )
            if archived_payload is None:
                raise ValueError(
                    "Historical option snapshots are unavailable from the live API and no archived "
                    f"raw payload exists for {logical_ts.isoformat()}."
                )
            payload = archived_payload
            raw_rows = 0

        normalized = normalize_option_chain_payload(payload, logical_ts=logical_ts, config=cfg)
        store.write(
            DatasetName.OPTION_CHAIN_SNAPSHOT,
            normalized,
            logical_ts=logical_ts,
            run_id=run_id,
            run_ts=started_at,
            write_mode=WriteMode.IMMUTABLE,
            stats={"rows": normalized.height, "source": "deribit_options_archiver"},
        )
        finished_at = deribit.now().astimezone(timezone.utc)
        duration_ms = int((finished_at - started_at).total_seconds() * 1000)
        _write_ingestion_run(
            store,
            logical_ts=logical_ts,
            run_id=run_id,
            run_ts=started_at,
            dataset=DatasetName.OPTION_CHAIN_SNAPSHOT.value,
            venue=venue,
            underlying=underlying,
            status=RunStatus.SUCCESS,
            rows_written_raw=raw_rows,
            rows_written_normalized=normalized.height,
            duration_ms=duration_ms,
        )
        return IngestionAttemptResult(
            logical_ts=logical_ts,
            run_id=run_id,
            status=RunStatus.SUCCESS,
            rows_written_raw=raw_rows,
            rows_written_normalized=normalized.height,
        )
    except Exception as exc:
        finished_at = deribit.now().astimezone(timezone.utc)
        duration_ms = int((finished_at - started_at).total_seconds() * 1000)
        _write_ingestion_run(
            store,
            logical_ts=logical_ts,
            run_id=run_id,
            run_ts=started_at,
            dataset=DatasetName.OPTION_CHAIN_SNAPSHOT.value,
            venue=venue,
            underlying=underlying,
            status=RunStatus.FAILED,
            rows_written_raw=0,
            rows_written_normalized=0,
            duration_ms=duration_ms,
            error_message=str(exc),
        )
        return IngestionAttemptResult(
            logical_ts=logical_ts,
            run_id=run_id,
            status=RunStatus.FAILED,
            rows_written_raw=0,
            rows_written_normalized=0,
            error_message=str(exc),
        )


def archive_gap(
    store: DatasetStore,
    *,
    logical_ts: datetime,
    venue: str = "deribit",
    underlying: str = "BTC",
) -> IngestionAttemptResult:
    run_ts = datetime.now(tz=timezone.utc)
    run_id = uuid4().hex
    _write_ingestion_run(
        store,
        logical_ts=logical_ts,
        run_id=run_id,
        run_ts=run_ts,
        dataset=DatasetName.OPTION_CHAIN_SNAPSHOT.value,
        venue=venue,
        underlying=underlying,
        status=RunStatus.GAP,
        rows_written_raw=0,
        rows_written_normalized=0,
        duration_ms=0,
        error_message="missing_live_snapshot",
    )
    return IngestionAttemptResult(
        logical_ts=logical_ts,
        run_id=run_id,
        status=RunStatus.GAP,
        rows_written_raw=0,
        rows_written_normalized=0,
        error_message="missing_live_snapshot",
    )


def process_archiver_cycle(
    store: DatasetStore,
    *,
    now: datetime,
    poll_underlying: str = "BTC",
    venue: str = "deribit",
    force_current_hour: bool = False,
    client: DeribitPublicClient | None = None,
    config: AppConfig | None = None,
) -> list[IngestionAttemptResult]:
    cfg = config or get_config()
    collectible = current_collectible_hour(now, cfg.ingestion.archive_grace_seconds)
    last_attempted = _latest_attempted_logical_ts(
        store,
        venue=venue,
        underlying=poll_underlying,
    )
    results: list[IngestionAttemptResult] = []
    if last_attempted is None:
        results.append(
            archive_deribit_options_once(
                store,
                as_of=collectible,
                force=force_current_hour,
                underlying=poll_underlying,
                venue=venue,
                client=client,
                config=cfg,
            )
        )
        return results

    next_expected = last_attempted + timedelta(hours=1)
    while next_expected < collectible:
        results.append(
            archive_gap(
                store,
                logical_ts=next_expected,
                venue=venue,
                underlying=poll_underlying,
            )
        )
        next_expected += timedelta(hours=1)

    if next_expected == collectible:
        results.append(
            archive_deribit_options_once(
                store,
                as_of=collectible,
                force=force_current_hour,
                underlying=poll_underlying,
                venue=venue,
                client=client,
                config=cfg,
            )
        )
    return results


def run_deribit_options_daemon(
    store: DatasetStore,
    *,
    poll_seconds: int,
    grace_seconds: int,
    force_current_hour: bool = False,
    underlying: str = "BTC",
    venue: str = "deribit",
    client: DeribitPublicClient | None = None,
    config: AppConfig | None = None,
    sleep_fn: Callable[[float], None] = time.sleep,
) -> None:
    cfg = config or get_config()
    effective_cfg = cfg.model_copy(
        update={
            "ingestion": cfg.ingestion.model_copy(
                update={
                    "archive_poll_seconds": poll_seconds,
                    "archive_grace_seconds": grace_seconds,
                }
            )
        }
    )
    deribit = client or DeribitPublicClient(effective_cfg)
    forced_once = False
    while True:
        results = process_archiver_cycle(
            store,
            now=deribit.now(),
            poll_underlying=underlying,
            venue=venue,
            force_current_hour=force_current_hour and not forced_once,
            client=deribit,
            config=effective_cfg,
        )
        if results and force_current_hour and not forced_once:
            forced_once = True
        for result in results:
            print(
                f"{result.logical_ts.isoformat()} status={result.status.value} "
                f"raw={result.rows_written_raw} normalized={result.rows_written_normalized} "
                f"run_id={result.run_id}",
                flush=True,
            )
        sleep_fn(float(poll_seconds))


def rebuild_option_chain_from_archive(
    store: DatasetStore,
    *,
    start: datetime,
    end: datetime,
    underlying: str = "BTC",
    venue: str = "deribit",
    config: AppConfig | None = None,
) -> dict[str, int]:
    cfg = config or get_config()
    archived = store.read(
        DatasetName.RAW_DERIBIT_OPTION_SNAPSHOT,
        start=start,
        end=end,
        filters={"venue": venue, "underlying": underlying},
    ).sort("logical_ts")
    if archived.is_empty():
        return {"snapshots": 0, "rows": 0, "partitions": 0}

    snapshots = 0
    rows = 0
    partitions = 0
    for raw in archived.iter_rows(named=True):
        logical_ts = _coerce_dt(raw["logical_ts"])
        payload = orjson.loads(str(raw["payload"]))
        frame = normalize_option_chain_payload(payload, logical_ts=logical_ts, config=cfg)
        if frame.is_empty():
            continue
        written = store.write(
            DatasetName.OPTION_CHAIN_SNAPSHOT,
            frame,
            logical_ts=logical_ts,
            run_id=f"rebuild-{uuid4().hex}",
            run_ts=datetime.now(tz=timezone.utc),
            write_mode=WriteMode.IMMUTABLE,
            stats={
                "rows": frame.height,
                "source": "archive_rebuild",
                "underlying": underlying,
            },
        )
        snapshots += 1
        rows += frame.height
        partitions += len(written)
    return {"snapshots": snapshots, "rows": rows, "partitions": partitions}


def _write_raw_snapshot(
    store: DatasetStore,
    *,
    payload: dict[str, Any],
    logical_ts: datetime,
    run_id: str,
    run_ts: datetime,
) -> int:
    observed_at = _coerce_dt(payload["observed_at"])
    raw_frame = pl.DataFrame(
        [
            {
                "ts": logical_ts,
                "date": logical_ts.date(),
                "venue": payload["venue"],
                "underlying": payload["underlying"],
                "payload": orjson.dumps(_serialize_payload(payload)).decode(),
                "instrument_count": len(payload.get("books", [])),
                "quality_flag": "ok",
                "logical_ts": logical_ts,
                "run_id": run_id,
                "run_ts": run_ts,
                "write_mode": WriteMode.IMMUTABLE.value,
            }
        ]
    )
    store.write(
        DatasetName.RAW_DERIBIT_OPTION_SNAPSHOT,
        raw_frame,
        logical_ts=logical_ts,
        run_id=run_id,
        run_ts=run_ts,
        write_mode=WriteMode.IMMUTABLE,
        stats={
            "rows": 1,
            "observed_at": observed_at,
            "instrument_count": len(payload.get("books", [])),
        },
    )
    return 1


def _write_ingestion_run(
    store: DatasetStore,
    *,
    logical_ts: datetime,
    run_id: str,
    run_ts: datetime,
    dataset: str,
    venue: str,
    underlying: str,
    status: RunStatus,
    rows_written_raw: int,
    rows_written_normalized: int,
    duration_ms: int,
    error_message: str | None = None,
) -> None:
    frame = pl.DataFrame(
        [
            {
                "ts": logical_ts,
                "date": logical_ts.date(),
                "dataset": dataset,
                "venue": venue,
                "underlying": underlying,
                "status": status.value,
                "rows_written_raw": rows_written_raw,
                "rows_written_normalized": rows_written_normalized,
                "error_message": error_message,
                "duration_ms": duration_ms,
                "logical_ts": logical_ts,
                "run_id": run_id,
                "run_ts": run_ts,
                "write_mode": WriteMode.IMMUTABLE.value,
            }
        ]
    )
    store.write(
        DatasetName.INGESTION_RUNS,
        frame,
        logical_ts=logical_ts,
        run_id=run_id,
        run_ts=run_ts,
        status=status,
        write_mode=WriteMode.IMMUTABLE,
        stats={
            "dataset": dataset,
            "rows_written_raw": rows_written_raw,
            "rows_written_normalized": rows_written_normalized,
        },
        error_reason=error_message,
    )


def _has_successful_snapshot(
    store: DatasetStore,
    logical_ts: datetime,
    *,
    venue: str,
    underlying: str,
) -> bool:
    runs = store.read(
        DatasetName.INGESTION_RUNS,
        start=logical_ts,
        end=logical_ts,
        filters={
            "dataset": DatasetName.OPTION_CHAIN_SNAPSHOT.value,
            "venue": venue,
            "underlying": underlying,
            "logical_ts": logical_ts,
            "status": RunStatus.SUCCESS.value,
        },
        read_mode=ReadMode.ALL_RUNS,
    )
    return not runs.is_empty()


def _load_archived_payload(
    store: DatasetStore,
    *,
    logical_ts: datetime,
    venue: str,
    underlying: str,
) -> dict[str, Any] | None:
    raw = store.read(
        DatasetName.RAW_DERIBIT_OPTION_SNAPSHOT,
        start=logical_ts,
        end=logical_ts,
        filters={"venue": venue, "underlying": underlying, "logical_ts": logical_ts},
        read_mode=ReadMode.LATEST,
    )
    if raw.is_empty():
        return None
    payload = orjson.loads(str(raw.item(0, "payload")))
    payload["observed_at"] = payload["observed_at"]
    return payload


def _latest_attempted_logical_ts(
    store: DatasetStore,
    *,
    venue: str,
    underlying: str,
) -> datetime | None:
    runs = store.read(
        DatasetName.INGESTION_RUNS,
        filters={
            "dataset": DatasetName.OPTION_CHAIN_SNAPSHOT.value,
            "venue": venue,
            "underlying": underlying,
        },
        read_mode=ReadMode.ALL_RUNS,
    )
    if runs.is_empty():
        return None
    return runs.sort("logical_ts").tail(1).item(0, "logical_ts")


def _serialize_payload(payload: dict[str, Any]) -> dict[str, Any]:
    serialized = dict(payload)
    serialized["observed_at"] = _coerce_dt(payload["observed_at"]).isoformat()
    return serialized


def _coerce_dt(value: Any) -> datetime:
    if isinstance(value, datetime):
        return value.astimezone(timezone.utc)
    if isinstance(value, str):
        return datetime.fromisoformat(value.replace("Z", "+00:00")).astimezone(timezone.utc)
    raise TypeError(f"Unsupported datetime value: {type(value)!r}")
