from __future__ import annotations

import asyncio
import hashlib
import logging
import os
import random
import time
from typing import Any, Dict, Optional, Tuple

import numpy as np
import pandas as pd
from fastapi import HTTPException

try:
    import torch
    TORCH_IMPORT_ERROR: Optional[str] = None
except Exception as exc:  # pragma: no cover - environment dependent
    torch = None  # type: ignore[assignment]
    TORCH_IMPORT_ERROR = str(exc)

KRONOS_AVAILABLE = False
KRONOS_IMPORT_ERROR: Optional[str] = None
Kronos = None
KronosTokenizer = None
KronosPredictor = None

if TORCH_IMPORT_ERROR:
    KRONOS_IMPORT_ERROR = TORCH_IMPORT_ERROR
else:
    try:
        from backend.kronos_core.model import (  # type: ignore[assignment]
            Kronos,
            KronosPredictor,
            KronosTokenizer,
        )

        KRONOS_AVAILABLE = True
    except Exception as exc:  # pragma: no cover - import environment dependent
        KRONOS_IMPORT_ERROR = str(exc)


class KronosForecaster:
    """Async OHLC4-only wrapper around the vendored Kronos model."""

    MODEL_HF_ID = os.getenv("KRONOS_MODEL_HF_ID", "NeoQuasar/Kronos-base")
    TOKENIZER_HF_ID = "NeoQuasar/Kronos-Tokenizer-base"
    MAX_CONTEXT = 512
    MIN_CONTEXT = 32
    DEFAULT_SAMPLE_COUNT = 5
    DEFAULT_TEMPERATURE = 0.9
    DEFAULT_TOP_P = 0.9
    DEFAULT_TOP_K = 0

    def __init__(self, logger: logging.Logger) -> None:
        self._logger = logger
        self._model: Optional[Any] = None
        self._tokenizer: Optional[Any] = None
        self._predictor: Optional[Any] = None
        self._loaded = False
        self._load_lock: Optional[asyncio.Lock] = None
        self._predict_lock: Optional[asyncio.Lock] = None

    async def _get_load_lock(self) -> asyncio.Lock:
        if self._load_lock is None:
            self._load_lock = asyncio.Lock()
        return self._load_lock

    async def _get_predict_lock(self) -> asyncio.Lock:
        if self._predict_lock is None:
            self._predict_lock = asyncio.Lock()
        return self._predict_lock

    @property
    def is_ready(self) -> bool:
        return self._loaded

    @property
    def device(self) -> str:
        if self._predictor is None:
            return "not_loaded"
        try:
            return str(self._predictor.device)
        except Exception:  # pragma: no cover - defensive
            return "cpu"

    @staticmethod
    def _extract_timestamps(df: pd.DataFrame) -> pd.Series:
        if "timestamps" in df.columns:
            ts = pd.to_datetime(df["timestamps"], utc=True)
        elif "time" in df.columns:
            ts = pd.to_datetime(df["time"], unit="s", utc=True)
        else:
            raise HTTPException(status_code=422, detail="Kronos requires timestamps or time column")
        if ts.isna().any():
            raise HTTPException(status_code=422, detail="Kronos input timestamps contain NaT")
        return ts.reset_index(drop=True)

    @classmethod
    def _build_ohlc4_proxy_frame(cls, df: pd.DataFrame) -> Tuple[pd.DataFrame, np.ndarray]:
        required_cols = {"open", "high", "low", "close"}
        missing_cols = sorted(required_cols - set(df.columns))
        if missing_cols:
            raise HTTPException(
                status_code=422,
                detail=f"Kronos input is missing OHLC columns: {', '.join(missing_cols)}",
            )

        ohlc4 = (
            df[["open", "high", "low", "close"]]
            .mean(axis=1)
            .to_numpy(dtype=np.float32, copy=False)
        )
        if ohlc4.ndim != 1 or len(ohlc4) < cls.MIN_CONTEXT:
            raise HTTPException(
                status_code=422,
                detail=f"Kronos requires at least {cls.MIN_CONTEXT} OHLC4 points",
            )
        if not np.isfinite(ohlc4).all():
            raise HTTPException(status_code=422, detail="Kronos input contains non-finite OHLC4 values")

        proxy = pd.DataFrame(
            {
                "open": ohlc4,
                "high": ohlc4,
                "low": ohlc4,
                "close": ohlc4,
                "volume": np.zeros(len(ohlc4), dtype=np.float32),
                "amount": np.zeros(len(ohlc4), dtype=np.float32),
            }
        )
        return proxy, ohlc4.astype(float)

    @staticmethod
    def _build_future_timestamps(
        history_timestamps: pd.Series,
        horizon: int,
        step_seconds: int,
    ) -> pd.Series:
        last_ts = pd.Timestamp(history_timestamps.iloc[-1])
        last_ts = last_ts.tz_convert("UTC") if last_ts.tzinfo is not None else last_ts.tz_localize("UTC")
        delta = pd.to_timedelta(step_seconds, unit="s")
        future_index = [last_ts + (delta * (step + 1)) for step in range(horizon)]
        return pd.Series(pd.DatetimeIndex(future_index), copy=False)

    @staticmethod
    def _stable_seed(
        symbol: str,
        interval: str,
        horizon: int,
        context_length: int,
        last_timestamp: pd.Timestamp,
    ) -> int:
        payload = "|".join(
            [
                symbol,
                interval,
                str(horizon),
                str(context_length),
                str(int(last_timestamp.timestamp())),
            ]
        )
        return int(hashlib.sha256(payload.encode("utf-8")).hexdigest()[:8], 16)

    @staticmethod
    def _build_output_validation(
        p10: np.ndarray,
        p50: np.ndarray,
        p90: np.ndarray,
        samples: np.ndarray,
    ) -> Dict[str, Any]:
        quantiles_monotonic = bool(
            np.all(p10 <= (p50 + 1e-6))
            and np.all(p50 <= (p90 + 1e-6))
        )
        if not quantiles_monotonic:
            raise HTTPException(status_code=500, detail="Kronos returned non-monotonic sample quantiles")
        if not (np.isfinite(p10).all() and np.isfinite(p50).all() and np.isfinite(p90).all()):
            raise HTTPException(status_code=500, detail="Kronos output contains non-finite values")

        return {
            "sample_shape": [int(dim) for dim in samples.shape],
            "quantile_shape": [1, int(p50.shape[0]), 3],
            "quantiles_monotonic": quantiles_monotonic,
            "quantile_source": "sample_paths",
        }

    async def _lazy_load(self) -> None:
        if self._loaded:
            return
        load_lock = await self._get_load_lock()
        async with load_lock:
            if self._loaded:
                return
            if not KRONOS_AVAILABLE or Kronos is None or KronosTokenizer is None or KronosPredictor is None:
                raise HTTPException(
                    status_code=503,
                    detail=(
                        "Kronos is unavailable. Ensure backend/kronos_core is present and "
                        "dependencies from its requirements are installed."
                    ),
                )
            try:
                self._logger.info("[Kronos] Loading %s ...", self.MODEL_HF_ID)
                tokenizer = await asyncio.to_thread(
                    KronosTokenizer.from_pretrained,
                    self.TOKENIZER_HF_ID,
                )
                model = await asyncio.to_thread(
                    Kronos.from_pretrained,
                    self.MODEL_HF_ID,
                )
                predictor = KronosPredictor(
                    model,
                    tokenizer,
                    device="cuda:0" if torch is not None and torch.cuda.is_available() else "cpu",
                    max_context=self.MAX_CONTEXT,
                )
                self._tokenizer = tokenizer
                self._model = model
                self._predictor = predictor
                self._loaded = True
                self._logger.info("[Kronos] Ready on %s", self.device)
            except Exception as exc:
                self._logger.error("[Kronos] Init failed: %s", exc, exc_info=True)
                raise HTTPException(status_code=500, detail=f"Kronos init failed: {exc}")

    async def forecast(
        self,
        df: pd.DataFrame,
        horizon: int,
        interval: str,
        step_seconds: int,
        symbol: str = "",
    ) -> Dict[str, Any]:
        await self._lazy_load()
        if self._predictor is None:
            raise HTTPException(status_code=500, detail="Kronos predictor is not initialized")

        history_timestamps = self._extract_timestamps(df)
        proxy_df, _ = self._build_ohlc4_proxy_frame(df)
        context_length = min(len(proxy_df), self.MAX_CONTEXT)
        proxy_context = proxy_df.tail(context_length).reset_index(drop=True)
        x_timestamp = history_timestamps.tail(context_length).reset_index(drop=True)
        y_timestamp = self._build_future_timestamps(
            history_timestamps=x_timestamp,
            horizon=horizon,
            step_seconds=step_seconds,
        )

        sample_count = self.DEFAULT_SAMPLE_COUNT
        seed = self._stable_seed(
            symbol=symbol or "UNKNOWN",
            interval=interval,
            horizon=horizon,
            context_length=context_length,
            last_timestamp=(
                pd.Timestamp(x_timestamp.iloc[-1]).tz_convert("UTC")
                if pd.Timestamp(x_timestamp.iloc[-1]).tzinfo is not None
                else pd.Timestamp(x_timestamp.iloc[-1]).tz_localize("UTC")
            ),
        )

        predict_lock = await self._get_predict_lock()
        async with predict_lock:
            try:
                random.seed(seed)
                np.random.seed(seed % (2**32 - 1))
                if torch is not None:
                    torch.manual_seed(seed)
                    if torch.cuda.is_available():
                        torch.cuda.manual_seed_all(seed)

                t0 = time.time()
                sample_tensor = await asyncio.to_thread(
                    self._predictor.predict_samples,
                    proxy_context,
                    x_timestamp,
                    y_timestamp,
                    horizon,
                    self.DEFAULT_TEMPERATURE,
                    self.DEFAULT_TOP_K,
                    self.DEFAULT_TOP_P,
                    sample_count,
                    False,
                )
                elapsed = time.time() - t0
                self._logger.info(
                    "[Kronos] %.2fs | horizon=%d ctx=%d samples=%d device=%s",
                    elapsed,
                    horizon,
                    context_length,
                    sample_count,
                    self.device,
                )
            except Exception as exc:
                self._logger.error("[Kronos] Forecast failed: %s", exc, exc_info=True)
                raise HTTPException(status_code=500, detail=f"Kronos prediction failed: {exc}")

        samples = np.asarray(sample_tensor, dtype=float)
        if samples.ndim != 3 or samples.shape[0] != sample_count or samples.shape[1] != horizon or samples.shape[2] < 4:
            raise HTTPException(
                status_code=500,
                detail=f"Kronos sample tensor has invalid shape: {list(samples.shape)}",
            )

        ohlc4_samples = samples[:, :, :4].mean(axis=2)
        p10 = np.quantile(ohlc4_samples, 0.1, axis=0).astype(float)
        p50 = np.quantile(ohlc4_samples, 0.5, axis=0).astype(float)
        p90 = np.quantile(ohlc4_samples, 0.9, axis=0).astype(float)
        output_validation = self._build_output_validation(p10, p50, p90, samples)

        return {
            "p10": p10,
            "p50": p50,
            "p90": p90,
            "model_name": self.MODEL_HF_ID,
            "tokenizer_name": self.TOKENIZER_HF_ID,
            "context_length": context_length,
            "output_horizon": horizon,
            "sample_count": sample_count,
            "seed": seed,
            "input_semantics": {
                "feature_channels": ["ohlc4"],
                "active_forecast_channels": ["ohlc4"],
                "proxy_channels": ["open", "high", "low", "close"],
                "ignored_channels": ["volume", "amount"],
                "price_mode": "ohlc4_single_channel",
                "base_signal": "ohlc4",
                "volume_mode": "synthetic_zero",
                "amount_mode": "synthetic_zero",
                "adapter_mode": "kronos_ohlc4_proxy",
            },
            "input_validation": {
                "series_field": "ohlc4",
                "dtype": "float32",
                "context_length": context_length,
                "finite": True,
            },
            "output_semantics": {
                "forecast_channel": "ohlc4",
                "forecast_mode": "single_future_ohlc4_line",
                "quantile_fields": ["p10", "p50", "p90"],
                "quantile_source": "sample_paths",
                "candle_projection": "omitted",
                "reference_baseline": "last_ohlc4",
            },
            "output_validation": output_validation,
        }