from __future__ import annotations

import json
import logging
import time
from pathlib import Path

logger = logging.getLogger(__name__)

_EXPECTED_KEY_PARTS = 2  # "SYMBOL:THRESHOLD" format


def _checkpoint_dir() -> Path:
    """Checkpoint directory for streaming state."""
    from platformdirs import user_cache_dir

    d = Path(user_cache_dir("opendeviationbar")) / "checkpoints"
    d.mkdir(parents=True, exist_ok=True)
    return d


def _checkpoint_path(symbol: str, threshold: int) -> Path:
    return _checkpoint_dir() / f"streaming_{symbol}_{threshold}.json"


def _load_checkpoint(symbol: str, threshold: int) -> dict | None:
    path = _checkpoint_path(symbol, threshold)
    if path.exists():
        data = json.loads(path.read_text())
        # Migrate legacy checkpoint keys (timestamp_ms → close_time_us → close_time_us)
        if "last_bar_timestamp_ms" in data and "last_bar_close_time_us" not in data:
            data["last_bar_close_time_us"] = data.pop("last_bar_timestamp_ms")
        if "last_bar_close_time_us" in data and "last_bar_close_time_us" not in data:
            data["last_bar_close_time_us"] = data.pop("last_bar_close_time_us")
            path.write_text(json.dumps(data))
        return data
    return None


def _save_checkpoint(
    symbol: str,
    threshold: int,
    bar_dict: dict,
    *,
    processor_checkpoint: dict | None = None,
) -> None:
    path = _checkpoint_path(symbol, threshold)
    checkpoint = {
        "symbol": symbol,
        "threshold": threshold,
        "last_bar_close_time_us": bar_dict.get("close_time_us"),
        "updated_at": time.time(),
    }
    # Issue #111 (Ariadne): Track trade ID at the streaming boundary
    last_trade_id = bar_dict.get("last_agg_trade_id")
    if last_trade_id and last_trade_id > 0:
        checkpoint["last_agg_trade_id"] = last_trade_id
    if processor_checkpoint is not None:
        checkpoint["processor_checkpoint"] = processor_checkpoint
    # Issue #97: Provenance metadata for forensic reconstruction
    from opendeviationbar.checkpoint import (
        _checkpoint_provenance,
    )

    checkpoint["provenance"] = _checkpoint_provenance()
    # Atomic write: temp + fsync + rename (Issue #97)
    tmp = path.with_suffix(".tmp")
    data = json.dumps(checkpoint)
    tmp.write_text(data)
    tmp.replace(path)


def _inject_checkpoints(
    engine: object,
    symbols: list[str],
    thresholds: list[int],
    trace_id: str,
    gap_fill_results: dict | None = None,
    valid_pairs: set[tuple[str, int]] | None = None,
) -> tuple[int, set[tuple[str, int]]]:
    """Inject saved processor checkpoints into engine before start.

    Parameters
    ----------
    gap_fill_results
        Dict keyed by ``"{symbol}@{threshold}"`` → ``BackfillResult`` from
        ``_gap_fill()``.  When gap-fill wrote ≥1 bar for a pair, the
        checkpoint for that pair is skipped — the gap-fill data is
        authoritative and the checkpoint's incomplete bar is obsolete
        (Issue #96).
    valid_pairs
        Issue #120: When provided, only inject checkpoints for these pairs.
        Eliminates useless checkpoint file lookups for invalid pairs.
    """
    from opendeviationbar.hooks import HookEvent, emit_hook
    from opendeviationbar.logging import log_checkpoint_event

    pairs = (
        valid_pairs
        if valid_pairs is not None
        else {(sym, thr) for sym in symbols for thr in thresholds}
    )

    loaded = 0
    loaded_pairs: set[tuple[str, int]] = set()
    for symbol, threshold in sorted(pairs):
        # Issue #96: Skip checkpoint when gap-fill already wrote bars
        key = f"{symbol}@{threshold}"
        gf = (gap_fill_results or {}).get(key)
        if gf is not None and getattr(gf, "bars_written", 0) > 0:
            logger.info(
                "skipping checkpoint for %s@%d: gap-fill wrote %d bars",
                symbol,
                threshold,
                gf.bars_written,
            )
            continue

        cp = _load_checkpoint(symbol, threshold)
        if not (cp and cp.get("processor_checkpoint")):
            continue

        # Guard: reject stale checkpoints that point beyond CH's last bar.
        cp_bar = cp["processor_checkpoint"].get("incomplete_bar", {})
        cp_first_tid = cp_bar.get("first_agg_trade_id", 0)
        try:
            from opendeviationbar.clickhouse.cache import OpenDeviationBarCache

            with OpenDeviationBarCache() as cache:
                _, ch_last_tid = cache.get_latest_bar_with_trade_id(symbol, threshold)
            if ch_last_tid and cp_first_tid > 0 and cp_first_tid > ch_last_tid + 1:
                logger.info(
                    "REJECTING stale checkpoint for %s@%d: checkpoint first_tid=%d "
                    "is beyond CH last_tid=%d — skipping (file preserved for next restart)",
                    symbol,
                    threshold,
                    cp_first_tid,
                    ch_last_tid,
                )
                continue
        except Exception:
            pass  # CH unavailable — inject checkpoint anyway (better than nothing)

        try:
            engine.set_checkpoint(
                symbol,
                threshold,
                cp["processor_checkpoint"],
            )
            loaded += 1
            loaded_pairs.add((symbol, threshold))
            cp_age = time.time() - cp.get("updated_at", time.time())
            logger.info(
                "injected checkpoint for %s@%d (age=%.0fs)",
                symbol,
                threshold,
                cp_age,
            )
            log_checkpoint_event(
                "sidecar_checkpoint_inject",
                symbol,
                trace_id,
                threshold=threshold,
                checkpoint_age_sec=round(cp_age),
                has_incomplete_bar=bool(
                    cp["processor_checkpoint"].get("has_incomplete_bar")
                ),
            )
            emit_hook(
                HookEvent.SIDECAR_CHECKPOINT_RESTORED,
                symbol=symbol,
                trace_id=trace_id,
                threshold_dbps=threshold,
                has_incomplete_bar=bool(
                    cp["processor_checkpoint"].get("has_incomplete_bar")
                ),
                checkpoint_age_sec=round(cp_age),
            )
        except Exception:
            logger.exception(
                "failed to inject checkpoint for %s@%d",
                symbol,
                threshold,
            )
            log_checkpoint_event(
                "sidecar_checkpoint_inject_failed",
                symbol,
                trace_id,
                threshold=threshold,
            )
    return loaded, loaded_pairs


def _extract_checkpoints(
    engine: object,
    trace_id: str,
    bars_written: int,
) -> None:
    """Extract and save processor checkpoints on shutdown."""
    from opendeviationbar.hooks import HookEvent, emit_hook
    from opendeviationbar.logging import log_checkpoint_event

    try:
        all_checkpoints = engine.collect_checkpoints(timeout_ms=5000)
        for key, cp_dict in all_checkpoints.items():
            parts = key.split(":")
            if len(parts) == _EXPECTED_KEY_PARTS:
                sym, thr_str = parts
                thr_int = int(thr_str)
                _save_checkpoint(
                    sym,
                    thr_int,
                    {},
                    processor_checkpoint=cp_dict,
                )
                log_checkpoint_event(
                    "sidecar_checkpoint_extract",
                    sym,
                    trace_id,
                    threshold=thr_int,
                    has_incomplete_bar=bool(cp_dict.get("has_incomplete_bar")),
                )
                emit_hook(
                    HookEvent.SIDECAR_CHECKPOINT_SAVED,
                    symbol=sym,
                    trace_id=trace_id,
                    threshold_dbps=thr_int,
                    has_incomplete_bar=bool(cp_dict.get("has_incomplete_bar")),
                    bars_written_session=bars_written,
                )
        logger.info(
            "saved %d processor checkpoints on shutdown",
            len(all_checkpoints),
        )
    except Exception:
        logger.exception("failed to collect/save processor checkpoints")
