from __future__ import annotations

import json
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, cast

import polars as pl

from uforia.storage import DatasetName, DatasetStore

from .builders import (
    build_microstructure_feature_series,
    build_microstructure_regime_series,
    build_microstructure_signal_series,
)
from .constants import MODEL_VERSION


def ingest_live_engine_state(
    store: DatasetStore,
    live_root: Path,
    *,
    assets: tuple[str, ...] = ("BTC", "ETH"),
    lookback_hours: int = 6,
    cursor_root: Path | None = None,
) -> dict[str, int]:
    start = datetime.now(tz=timezone.utc) - timedelta(hours=lookback_hours)
    summary = {"book_states": 0, "market_events": 0, "funding_rows": 0}
    book_cursor = _load_cursor_state(cursor_root / "book_state.json") if cursor_root else None
    market_cursor = _load_cursor_state(cursor_root / "market.json") if cursor_root else None
    funding_cursor = _load_cursor_state(cursor_root / "funding.json") if cursor_root else None
    summary["book_states"] = _write_chunked_rows(
        store,
        DatasetName.MICROSTRUCTURE_BOOK_STATE,
        _iter_book_state_rows(
            live_root,
            start=start,
            assets=assets,
            cursor_state=book_cursor,
        ),
    )
    summary["market_events"] = _write_chunked_rows(
        store,
        DatasetName.MICROSTRUCTURE_RAW_EVENT,
        _iter_market_event_rows(
            live_root,
            start=start,
            assets=assets,
            cursor_state=market_cursor,
        ),
    )
    summary["funding_rows"] = _write_chunked_rows(
        store,
        DatasetName.FUNDING_RATE_SNAPSHOT,
        _iter_funding_rows(
            live_root,
            start=start,
            assets=assets,
            cursor_state=funding_cursor,
        ),
        chunk_size=512,
    )
    if cursor_root:
        _save_cursor_state(cursor_root / "book_state.json", book_cursor)
        _save_cursor_state(cursor_root / "market.json", market_cursor)
        _save_cursor_state(cursor_root / "funding.json", funding_cursor)
    return summary


def build_live_microstructure_bars(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    *,
    asset: str,
    bar_sizes: tuple[int, ...] = (5, 60),
) -> pl.DataFrame:
    book = store.read(
        DatasetName.MICROSTRUCTURE_BOOK_STATE,
        start=start,
        end=end,
        filters={"asset": asset},
    ).sort("ts")
    raw = store.read(
        DatasetName.MICROSTRUCTURE_RAW_EVENT,
        start=start,
        end=end,
        filters={"asset": asset, "event_type": "trade"},
    ).sort("ts")
    if book.is_empty():
        return pl.DataFrame()

    rows: list[pl.DataFrame] = []
    for bar_size in bar_sizes:
        every = f"{bar_size}s"
        trade_agg = (
            raw.group_by_dynamic(
                "ts",
                every=every,
                group_by=["asset", "venue"],
                closed="left",
            )
            .agg(
                [
                    pl.col("price").last().alias("trade_last"),
                    pl.col("price").mean().alias("vwap"),
                    pl.col("size").sum().alias("volume"),
                    pl.when(pl.col("side") == "buy")
                    .then(pl.col("size"))
                    .otherwise(0.0)
                    .sum()
                    .alias("buy_volume"),
                    pl.when(pl.col("side") == "sell")
                    .then(pl.col("size"))
                    .otherwise(0.0)
                    .sum()
                    .alias("sell_volume"),
                    pl.len().alias("trade_count"),
                ]
            )
            .rename({"ts": "bucket_ts"})
        )
        book_agg = (
            book.group_by_dynamic(
                "ts",
                every=every,
                group_by=["asset", "venue"],
                closed="left",
            )
            .agg(
                [
                    pl.col("date").first().alias("date"),
                    pl.col("mid_price").first().alias("open"),
                    pl.col("mid_price").max().alias("high"),
                    pl.col("mid_price").min().alias("low"),
                    pl.col("mid_price").last().alias("close"),
                    pl.col("microprice").last().alias("microprice"),
                    pl.col("spread").mean().alias("spread"),
                    pl.col("depth_bid_5").mean().alias("depth_bid_5"),
                    pl.col("depth_ask_5").mean().alias("depth_ask_5"),
                    pl.col("is_gap_detected").max().alias("is_gap"),
                ]
            )
            .rename({"ts": "bucket_ts"})
        )
        merged = (
            book_agg.join(trade_agg, on=["bucket_ts", "asset", "venue"], how="left")
            .with_columns(
                [
                    pl.col("bucket_ts").alias("ts"),
                    pl.lit(None, dtype=pl.String).alias("exchange_symbol"),
                    _venue_group_expr("venue").alias("venue_group"),
                    _market_type_expr("venue").alias("market_type"),
                    (_market_type_expr("venue") != pl.lit("spot")).alias("is_derivatives"),
                    _contract_type_expr("venue").alias("contract_type"),
                    pl.col("volume").fill_null(0.0),
                    pl.col("buy_volume").fill_null(0.0),
                    pl.col("sell_volume").fill_null(0.0),
                    pl.col("trade_count").fill_null(0),
                    pl.coalesce("vwap", "close").alias("vwap"),
                    pl.col("is_gap").fill_null(False),
                    pl.lit("event_time").alias("clock"),
                    pl.lit(bar_size).cast(pl.Int64).alias("bar_size_secs"),
                    (pl.col("depth_bid_5") * pl.col("close")).alias("depth_notional_bid_5"),
                    (pl.col("depth_ask_5") * pl.col("close")).alias("depth_notional_ask_5"),
                    pl.when(pl.col("trade_count") > 0)
                    .then(pl.lit("ok"))
                    .otherwise(pl.lit("book_only"))
                    .alias("quality_flag"),
                ]
            )
            .select(
                [
                    "ts",
                    "date",
                    "asset",
                    "venue",
                    "exchange_symbol",
                    "venue_group",
                    "market_type",
                    "is_derivatives",
                    "contract_type",
                    "clock",
                    "bar_size_secs",
                    "open",
                    "high",
                    "low",
                    "close",
                    "volume",
                    "buy_volume",
                    "sell_volume",
                    pl.col("trade_count").cast(pl.Int64),
                    "vwap",
                    "microprice",
                    "spread",
                    "depth_bid_5",
                    "depth_ask_5",
                    "depth_notional_bid_5",
                    "depth_notional_ask_5",
                    "is_gap",
                    "quality_flag",
                ]
            )
        )
        rows.append(merged)
    return pl.concat(rows, how="vertical_relaxed") if rows else pl.DataFrame()


def run_live_microstructure_inference_once(
    store: DatasetStore,
    live_root: Path | None = None,
    *,
    assets: tuple[str, ...] = ("BTC", "ETH"),
    horizons: tuple[str, ...] = ("30s", "1m", "2m", "5m"),
    lookback_hours: int = 6,
    cursor_root: Path | None = None,
) -> dict[str, Any]:
    from .training import generate_microstructure_model_output

    now = datetime.now(tz=timezone.utc)
    start = now - timedelta(hours=lookback_hours)
    ingest_summary = {"book_states": 0, "market_events": 0, "funding_rows": 0}
    if live_root is not None:
        ingest_summary = ingest_live_engine_state(
            store,
            live_root,
            assets=assets,
            lookback_hours=lookback_hours,
            cursor_root=cursor_root,
        )
    bar_paths = 0
    feature_paths = 0
    regime_paths = 0
    output_paths = 0
    signal_paths = 0
    for asset in assets:
        bars = build_live_microstructure_bars(store, start, now, asset=asset)
        if not bars.is_empty():
            bar_paths += len(store.write(DatasetName.MICROSTRUCTURE_BAR, bars))
        features = build_microstructure_feature_series(
            store,
            start,
            now,
            asset=asset,
            horizons=horizons,
        )
        if not features.is_empty():
            feature_paths += len(store.write(DatasetName.MICROSTRUCTURE_FEATURE_SERIES, features))
        regime = build_microstructure_regime_series(store, start, now, asset=asset)
        if not regime.is_empty():
            regime_paths += len(store.write(DatasetName.MICROSTRUCTURE_REGIME_SERIES, regime))

        for horizon in horizons:
            latest_run = (
                store.read(
                    DatasetName.MICROSTRUCTURE_TRAINING_RUN,
                    filters={
                        "asset": asset,
                        "horizon": horizon,
                        "source": "tcn_boost_fusion",
                        "model_version": MODEL_VERSION,
                    },
                )
                .sort("ts")
                .tail(1)
            )
            if latest_run.is_empty():
                continue
            artifact_payload = latest_run.item(0, "artifact_payload")
            if artifact_payload in {None, ""}:
                continue
            try:
                output_frame, signal_frame, _ = generate_microstructure_model_output(
                    store,
                    json.loads(str(artifact_payload)),
                    asset=asset,
                    horizon=horizon,
                )
                if not output_frame.is_empty():
                    output_paths += len(
                        store.write(DatasetName.MICROSTRUCTURE_MODEL_OUTPUT, output_frame)
                    )
                if not signal_frame.is_empty():
                    signal_paths += len(
                        store.write(DatasetName.MICROSTRUCTURE_SIGNAL_SERIES, signal_frame)
                    )
            except Exception as exc:
                print(
                    json.dumps(
                        {
                            "event": "microstructure_model_output_error",
                            "asset": asset,
                            "horizon": horizon,
                            "error": str(exc),
                        }
                    ),
                    flush=True,
                )
        paper = build_microstructure_signal_series(
            store,
            start,
            now,
            asset=asset,
            horizons=horizons,
        )
        if not paper.is_empty():
            signal_paths += len(store.write(DatasetName.MICROSTRUCTURE_SIGNAL_SERIES, paper))
    return {
        "ingest": ingest_summary,
        "bar_paths": bar_paths,
        "feature_paths": feature_paths,
        "regime_paths": regime_paths,
        "output_paths": output_paths,
        "signal_paths": signal_paths,
    }


def run_live_microstructure_inference_daemon(
    store: DatasetStore,
    live_root: Path | None = None,
    *,
    assets: tuple[str, ...] = ("BTC", "ETH"),
    horizons: tuple[str, ...] = ("30s", "1m", "2m", "5m"),
    interval_seconds: int = 15,
    lookback_hours: int = 6,
) -> None:
    cursor_root = live_root / ".ingest_cursors" if live_root is not None else None
    while True:
        try:
            result = run_live_microstructure_inference_once(
                store,
                live_root,
                assets=assets,
                horizons=horizons,
                lookback_hours=lookback_hours,
                cursor_root=cursor_root,
            )
            print(json.dumps(result), flush=True)
        except Exception as exc:
            print(
                json.dumps(
                    {
                        "event": "microstructure_inference_daemon_error",
                        "error": str(exc),
                    }
                ),
                flush=True,
            )
        time.sleep(interval_seconds)


def _read_book_state(
    live_root: Path,
    *,
    start: datetime,
    assets: tuple[str, ...],
    cursor_state: dict[str, int] | None = None,
) -> pl.DataFrame:
    rows = list(
        _iter_book_state_rows(
            live_root,
            start=start,
            assets=assets,
            cursor_state=cursor_state,
        )
    )
    return pl.DataFrame(rows, infer_schema_length=None) if rows else pl.DataFrame()


def _read_market_events(
    live_root: Path,
    *,
    start: datetime,
    assets: tuple[str, ...],
    cursor_state: dict[str, int] | None = None,
) -> pl.DataFrame:
    rows = list(
        _iter_market_event_rows(
            live_root,
            start=start,
            assets=assets,
            cursor_state=cursor_state,
        )
    )
    return pl.DataFrame(rows, infer_schema_length=None) if rows else pl.DataFrame()


def _read_funding_events(
    live_root: Path,
    *,
    start: datetime,
    assets: tuple[str, ...],
    cursor_state: dict[str, int] | None = None,
) -> pl.DataFrame:
    rows = list(
        _iter_funding_rows(
            live_root,
            start=start,
            assets=assets,
            cursor_state=cursor_state,
        )
    )
    return pl.DataFrame(rows, infer_schema_length=None) if rows else pl.DataFrame()


def _iter_book_state_rows(
    live_root: Path,
    *,
    start: datetime,
    assets: tuple[str, ...],
    cursor_state: dict[str, int] | None = None,
):
    for line in _iter_jsonl_partitioned(live_root / "book_state", start, cursor_state=cursor_state):
        payload = _parse_json(line)
        if not payload:
            continue
        asset = str(payload.get("asset", "")).upper()
        if asset not in assets:
            continue
        ts = _parse_dt(payload.get("timestamp"))
        if ts is None or ts < start:
            continue
        yield {
            "ts": ts,
            "date": ts.date(),
            "asset": asset,
            "venue": _normalize_venue(payload.get("venue")),
            "exchange_symbol": str(
                payload.get("symbol")
                or payload.get("exchange_symbol")
                or payload.get("venue")
                or ""
            ),
            "venue_group": _venue_group(payload.get("venue")),
            "market_type": _market_type(payload.get("venue")),
            "is_derivatives": _market_type(payload.get("venue")) != "spot",
            "contract_type": _contract_type(payload.get("venue")),
            "best_bid": payload.get("best_bid_price"),
            "best_ask": payload.get("best_ask_price"),
            "best_bid_size": payload.get("best_bid_size"),
            "best_ask_size": payload.get("best_ask_size"),
            "mid_price": payload.get("mid"),
            "microprice": payload.get("microprice"),
            "spread": payload.get("spread"),
            "spread_rel": payload.get("spread_rel"),
            "depth_bid_5": payload.get("depth_bid_5"),
            "depth_ask_5": payload.get("depth_ask_5"),
            "depth_notional_bid_5": _mul_or_none(
                _as_float(payload.get("depth_bid_5")),
                _as_float(payload.get("mid")),
            ),
            "depth_notional_ask_5": _mul_or_none(
                _as_float(payload.get("depth_ask_5")),
                _as_float(payload.get("mid")),
            ),
            "imbalance_5": payload.get("imbalance_5"),
            "is_seeded": payload.get("is_seeded", payload.get("is_complete")),
            "is_sequence_valid": payload.get("is_sequence_valid", payload.get("is_complete")),
            "is_gap_detected": payload.get("is_gap_detected", False),
            "reseed_count": payload.get("reseed_count", 0),
            "book_quality": payload.get(
                "book_quality",
                "ok" if payload.get("is_complete") else "partial",
            ),
            "quality_flag": "ok" if payload.get("is_complete") else "partial",
        }


def _iter_market_event_rows(
    live_root: Path,
    *,
    start: datetime,
    assets: tuple[str, ...],
    cursor_state: dict[str, int] | None = None,
):
    for line in _iter_jsonl_partitioned(
        live_root / "raw" / "market",
        start,
        cursor_state=cursor_state,
    ):
        wrapper = _parse_json(line)
        if not wrapper:
            continue
        payload = cast(dict[str, Any], wrapper.get("payload", {}))
        event_type = str(wrapper.get("event_type", ""))
        asset = str(payload.get("asset", "")).upper()
        if asset not in assets:
            continue
        ts = _parse_dt(payload.get("timestamp"))
        if ts is None or ts < start:
            continue
        yield {
            "ts": ts,
            "date": ts.date(),
            "asset": asset,
            "venue": _normalize_venue(payload.get("venue")),
            "event_type": event_type,
            "event_id": str(
                payload.get("trade_id")
                or payload.get("last_update_id")
                or f"{event_type}-{ts.isoformat()}"
            ),
            "exchange_symbol": str(
                payload.get("symbol")
                or payload.get("exchange_symbol")
                or payload.get("venue")
                or ""
            ),
            "venue_group": _venue_group(payload.get("venue")),
            "market_type": _market_type(payload.get("venue")),
            "is_derivatives": _market_type(payload.get("venue")) != "spot",
            "contract_type": _contract_type(payload.get("venue")),
            "exchange_ts": ts,
            "received_ts": _parse_dt(wrapper.get("received_at")),
            "local_ts": _parse_dt(wrapper.get("received_at")),
            "transaction_ts": ts,
            "side": _normalize_side(payload.get("side")),
            "liquidation_side": _normalize_side(
                payload.get("liquidation_side") or payload.get("liquidationSide")
            ),
            "price": _as_float(payload.get("price")),
            "size": _as_float(payload.get("size")),
            "trade_notional": _mul_or_none(
                _as_float(payload.get("price")),
                _as_float(payload.get("size")),
            ),
            "bid": _as_float(payload.get("bid_price")),
            "ask": _as_float(payload.get("ask_price")),
            "mark_price": _as_float(payload.get("mark_price")),
            "open_interest": _as_float(payload.get("oi_value")),
            "depth_notional_bid_5": _mul_or_none(
                _as_float(payload.get("depth_bid_5")),
                _as_float(payload.get("price")),
            ),
            "depth_notional_ask_5": _mul_or_none(
                _as_float(payload.get("depth_ask_5")),
                _as_float(payload.get("price")),
            ),
            "first_update_id": _as_int(payload.get("first_update_id")),
            "last_update_id": _as_int(payload.get("last_update_id")),
            "payload": json.dumps(payload, separators=(",", ":"), sort_keys=True),
            "quality_flag": "ok",
        }


def _iter_funding_rows(
    live_root: Path,
    *,
    start: datetime,
    assets: tuple[str, ...],
    cursor_state: dict[str, int] | None = None,
):
    for line in _iter_jsonl_partitioned(
        live_root / "raw" / "funding",
        start,
        cursor_state=cursor_state,
    ):
        wrapper = _parse_json(line)
        if not wrapper:
            continue
        payload = cast(dict[str, Any], wrapper.get("payload", {}))
        asset = str(payload.get("asset", "")).upper()
        if asset not in assets:
            continue
        ts = _parse_dt(payload.get("timestamp"))
        if ts is None or ts < start:
            continue
        if wrapper.get("event_type") != "funding":
            continue
        yield {
            "ts": ts,
            "date": ts.date(),
            "venue": _normalize_venue(payload.get("venue")),
            "underlying": asset,
            "source": "microstructure_live",
            "interest_1h": payload.get("rate"),
            "interest_8h": None,
            "quality_flag": "ok",
        }


def _write_chunked_rows(
    store: DatasetStore,
    dataset: DatasetName,
    rows,
    *,
    chunk_size: int = 10_000,
) -> int:
    buffer: list[dict[str, Any]] = []
    written = 0
    for row in rows:
        buffer.append(row)
        if len(buffer) >= chunk_size:
            written += len(store.write(dataset, pl.DataFrame(buffer, infer_schema_length=None)))
            buffer.clear()
    if buffer:
        written += len(store.write(dataset, pl.DataFrame(buffer, infer_schema_length=None)))
    return written


def _read_jsonl_partitioned(root: Path, start: datetime) -> list[str]:
    return list(_iter_jsonl_partitioned(root, start))


def _iter_jsonl_partitioned(
    root: Path,
    start: datetime,
    *,
    cursor_state: dict[str, int] | None = None,
):
    if not root.exists():
        return
    current_day = start.date()
    end_day = datetime.now(tz=timezone.utc).date()
    active_paths: set[str] = set()
    while current_day <= end_day:
        day = current_day.isoformat()
        file_path = root / f"date={day}" / "events.jsonl"
        if "book_state" in str(root):
            file_path = root / f"date={day}" / "state.jsonl"
        active_paths.add(str(file_path))
        if not file_path.exists():
            current_day += timedelta(days=1)
            continue
        with file_path.open("rb") as handle:
            if cursor_state is not None:
                size = file_path.stat().st_size
                offset = int(cursor_state.get(str(file_path), 0))
                if offset < 0 or offset > size:
                    offset = 0
                handle.seek(offset)
            for raw_line in handle:
                if cursor_state is not None:
                    cursor_state[str(file_path)] = handle.tell()
                line = raw_line.decode("utf-8", errors="ignore").strip()
                if line:
                    yield line
        current_day += timedelta(days=1)
    if cursor_state is not None:
        for path in list(cursor_state):
            if path not in active_paths and not Path(path).exists():
                cursor_state.pop(path, None)


def _load_cursor_state(path: Path) -> dict[str, int]:
    if not path.exists():
        return {}
    try:
        payload = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError):
        return {}
    if not isinstance(payload, dict):
        return {}
    state: dict[str, int] = {}
    for key, value in payload.items():
        if isinstance(key, str) and isinstance(value, int):
            state[key] = value
    return state


def _save_cursor_state(path: Path, state: dict[str, int] | None) -> None:
    if state is None:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(state, separators=(",", ":"), sort_keys=True))


def _parse_json(line: str) -> dict[str, Any] | None:
    if not line.strip():
        return None
    try:
        parsed = json.loads(line)
    except json.JSONDecodeError:
        return None
    return parsed if isinstance(parsed, dict) else None


def _parse_dt(value: Any) -> datetime | None:
    if not value:
        return None
    if isinstance(value, datetime):
        return value.astimezone(timezone.utc)
    text = str(value).replace("Z", "+00:00")
    try:
        parsed = datetime.fromisoformat(text)
    except ValueError:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


def _normalize_venue(value: Any) -> str:
    text = str(value or "")
    return {
        "BinanceSpot": "binance_spot",
        "BinanceFutures": "binance_futures",
        "BybitSpot": "bybit_spot",
        "BybitPerps": "bybit_perps",
        "Deribit": "deribit",
    }.get(text, text.lower())


def _venue_group(value: Any) -> str:
    venue = _normalize_venue(value)
    if venue.startswith("binance"):
        return "binance"
    if venue.startswith("bybit"):
        return "bybit"
    if venue.startswith("deribit"):
        return "deribit"
    return "other"


def _market_type(value: Any) -> str:
    venue = _normalize_venue(value)
    if "spot" in venue:
        return "spot"
    if "futures" in venue:
        return "futures"
    if "perps" in venue:
        return "perp"
    return "perp"


def _contract_type(value: Any) -> str:
    market_type = _market_type(value)
    if market_type == "spot":
        return "spot"
    if market_type == "futures":
        return "linear_future"
    return "linear_perp"


def _venue_group_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("binance"))
        .then(pl.lit("binance"))
        .when(pl.col(column).str.contains("bybit"))
        .then(pl.lit("bybit"))
        .when(pl.col(column).str.contains("deribit"))
        .then(pl.lit("deribit"))
        .otherwise(pl.lit("other"))
    )


def _market_type_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("spot"))
        .then(pl.lit("spot"))
        .when(pl.col(column).str.contains("futures"))
        .then(pl.lit("futures"))
        .when(pl.col(column).str.contains("perps"))
        .then(pl.lit("perp"))
        .otherwise(pl.lit("perp"))
    )


def _contract_type_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("spot"))
        .then(pl.lit("spot"))
        .when(pl.col(column).str.contains("futures"))
        .then(pl.lit("linear_future"))
        .otherwise(pl.lit("linear_perp"))
    )


def _normalize_side(value: Any) -> str | None:
    if value is None:
        return None
    text = str(value)
    if text in {"Buy", "buy"}:
        return "buy"
    if text in {"Sell", "sell"}:
        return "sell"
    return text.lower()


def _as_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _as_int(value: Any) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _mul_or_none(left: float | None, right: float | None) -> float | None:
    if left is None or right is None:
        return None
    return left * right
