"""
Kronos Financial Foundation Model — HuggingFace Spaces FastAPI server.

API:
  POST /api/load-model   {"model_key": "kronos-small"}
  POST /api/predict      {"k_lines": [[ts_ms,o,h,l,c,v],...],
                          "prediction_params": {"pred_len": 24}}

Auth: Authorization: Bearer <KRONOS_BEARER_TOKEN>
      Set KRONOS_BEARER_TOKEN as a HF Space secret (Settings → Variables and secrets).
"""
from __future__ import annotations

import os
import sys
import time
from contextlib import asynccontextmanager
from typing import Any

import numpy as np
import pandas as pd
import torch

# Kronos ships its own model.py — not available on PyPI
sys.path.insert(0, "/app/kronos")
from model import Kronos, KronosTokenizer, KronosPredictor  # noqa: E402

from fastapi import Depends, FastAPI, HTTPException, Security, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

BEARER_TOKEN: str = os.getenv("KRONOS_BEARER_TOKEN", "")

MODEL_REPO: dict[str, str] = {
    "kronos-small": "NeoQuasar/Kronos-small",
    "kronos-mini":  "NeoQuasar/Kronos-mini",
    "kronos-base":  "NeoQuasar/Kronos-base",
}
TOKENIZER_REPO = "NeoQuasar/Kronos-Tokenizer-base"
DEFAULT_MODEL_KEY = "kronos-small"

# ---------------------------------------------------------------------------
# Global model state
# ---------------------------------------------------------------------------

_predictor: KronosPredictor | None = None
_loaded_model_key: str | None = None

# ---------------------------------------------------------------------------
# Auth
# ---------------------------------------------------------------------------

_security = HTTPBearer(auto_error=False)


def _verify_token(
    creds: HTTPAuthorizationCredentials | None = Security(_security),
) -> None:
    if not BEARER_TOKEN:
        return
    if creds is None or creds.credentials != BEARER_TOKEN:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing Bearer token",
            headers={"WWW-Authenticate": "Bearer"},
        )


# ---------------------------------------------------------------------------
# Lifespan
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    try:
        _load_model(DEFAULT_MODEL_KEY)
    except Exception as exc:
        print(f"WARNING: could not pre-warm model on startup: {exc}")
    yield


app = FastAPI(title="Kronos API", lifespan=lifespan)

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

def _load_model(model_key: str) -> None:
    global _predictor, _loaded_model_key
    if _loaded_model_key == model_key and _predictor is not None:
        return
    repo_id = MODEL_REPO.get(model_key)
    if repo_id is None:
        raise ValueError(f"Unknown model_key: {model_key!r}. Valid keys: {list(MODEL_REPO)}")
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    tokenizer = KronosTokenizer.from_pretrained(TOKENIZER_REPO)
    model = Kronos.from_pretrained(repo_id)
    _predictor = KronosPredictor(model, tokenizer, device=device, max_context=512)
    _loaded_model_key = model_key


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

def _run_kronos(k_lines: list[list[float]], pred_len: int) -> list[list[float]]:
    """
    Input:  [[ts_ms, open, high, low, close, volume], ...]
    Output: [[ts_ms, open, high, low, close], ...]
    """
    if _predictor is None:
        raise RuntimeError("Model not loaded.")

    interval_ms = float(k_lines[1][0] - k_lines[0][0]) if len(k_lines) > 1 else 3_600_000.0
    last_ts_ms = float(k_lines[-1][0])

    x_timestamps = pd.to_datetime([row[0] for row in k_lines], unit="ms", utc=True)
    y_timestamps = pd.to_datetime(
        [last_ts_ms + (i + 1) * interval_ms for i in range(pred_len)],
        unit="ms",
        utc=True,
    )

    df = pd.DataFrame(k_lines, columns=["ts_ms", "open", "high", "low", "close", "volume"])
    df["amount"] = df["close"] * df["volume"]
    x_df = df[["open", "high", "low", "close", "volume", "amount"]]

    pred_df = _predictor.predict(
        df=x_df,
        x_timestamp=x_timestamps,
        y_timestamp=y_timestamps,
        pred_len=pred_len,
        T=1.0,
        top_p=0.9,
        sample_count=1,
    )

    result: list[list[float]] = []
    for i, row in enumerate(pred_df.itertuples(index=False)):
        ts = last_ts_ms + (i + 1) * interval_ms
        o = float(getattr(row, "open", row[0]))
        h = float(getattr(row, "high", row[1]))
        l = float(getattr(row, "low", row[2]))
        c = float(getattr(row, "close", row[3]))
        result.append([ts, o, h, l, c])

    return result


# ---------------------------------------------------------------------------
# Request / response schemas
# ---------------------------------------------------------------------------

class LoadModelRequest(BaseModel):
    model_key: str = DEFAULT_MODEL_KEY


class PredictRequest(BaseModel):
    k_lines: list[list[float]]
    prediction_params: dict[str, Any] = {}


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get("/health")
def health():
    return {"status": "ok", "model_loaded": _loaded_model_key}


@app.post("/api/load-model", dependencies=[Depends(_verify_token)])
def load_model(req: LoadModelRequest):
    try:
        t0 = time.time()
        _load_model(req.model_key)
        return {
            "status": "loaded",
            "model_key": req.model_key,
            "load_time_s": round(time.time() - t0, 2),
        }
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/api/predict", dependencies=[Depends(_verify_token)])
def predict(req: PredictRequest):
    if _predictor is None:
        raise HTTPException(status_code=400, detail="Model not loaded. POST /api/load-model first.")
    pred_len: int = int(req.prediction_params.get("pred_len", 24))
    try:
        prediction_results = _run_kronos(req.k_lines, pred_len)
        return {"prediction_results": prediction_results}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))