from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any, Iterator, List, Optional

import httpx
import numpy as np
from gluonts.dataset import Dataset as GluonDataset
from gluonts.model import Forecast
from gluonts.model.forecast import QuantileForecast
from gluonts.model.predictor import RepresentablePredictor

from tsfm_bench.eval.predictors import _ForecastConfig

DEFAULT_API_URL = "https://api.tsfm.ai/v1/forecast"
DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]


@dataclass
class TsfmApiConfig:
    model_id: str
    api_url: str = DEFAULT_API_URL
    api_key: str | None = None
    timeout: float = 120.0
    freq: str = "H"


class TsfmApiPredictor(RepresentablePredictor):
    """Zero-shot forecaster backed by the TSFM.ai hosted API."""

    def __init__(
        self,
        config: TsfmApiConfig,
        prediction_length: int,
        quantile_levels: Optional[List[float]] = None,
    ):
        super().__init__(prediction_length=prediction_length)
        self.config = config
        self.quantile_levels = quantile_levels or DEFAULT_QUANTILES
        self.forecast_config = _ForecastConfig.from_quantiles(self.quantile_levels)
        self._api_key = config.api_key or os.getenv("TSFM_API_KEY")
        self.leaderboard_name = config.model_id.split("/")[-1]

    def predict(self, dataset: GluonDataset, **kwargs) -> Iterator[Forecast]:
        if not self._api_key:
            raise RuntimeError(
                "TSFM_API_KEY is required for online zero-shot evaluation. "
                "Get a key at https://tsfm.ai/"
            )

        headers = {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

        with httpx.Client(timeout=self.config.timeout, trust_env=False) as client:
            for entry in dataset:
                target = np.asarray(entry["target"], dtype=np.float64)
                univariate = target.ndim == 1
                working = target.reshape(1, -1) if univariate else target

                forecast_rows = []
                for variate_idx in range(working.shape[0]):
                    series = working[variate_idx]
                    series = series[np.isfinite(series)]
                    payload = {
                        "model": self.config.model_id,
                        "inputs": [
                            {
                                "item_id": str(entry.get("item_id", variate_idx)),
                                "target": _to_api_target(series),
                                "start": _format_start(entry.get("start")),
                            }
                        ],
                        "parameters": {
                            "prediction_length": self.prediction_length,
                            "freq": self.config.freq,
                            "quantiles": self.quantile_levels,
                        },
                    }
                    response = client.post(
                        self.config.api_url,
                        headers=headers,
                        json=payload,
                    )
                    response.raise_for_status()
                    forecast_rows.append(
                        _parse_api_forecast(
                            response.json(),
                            self.prediction_length,
                            self.quantile_levels,
                        )
                    )

                stacked = np.stack(forecast_rows, axis=0)
                if univariate:
                    forecast_arrays = stacked[0]
                else:
                    forecast_arrays = stacked

                yield QuantileForecast(
                    forecast_arrays=forecast_arrays,
                    forecast_keys=self.forecast_config.forecast_keys,
                    start_date=entry["start"] + len(entry["target"]),
                    item_id=entry["item_id"],
                )


def _to_api_target(series: np.ndarray) -> list[list[float]]:
    """TSFM.ai expects target shaped [num_timesteps][num_channels]."""
    return [[float(v)] for v in series.tolist()]


def _format_start(start: Any) -> str:
    if hasattr(start, "to_timestamp"):
        return start.to_timestamp().isoformat()
    return str(start)


def _parse_api_forecast(
    payload: dict[str, Any],
    prediction_length: int,
    quantile_levels: list[float],
) -> np.ndarray:
    """Parse TSFM.ai response into (num_outputs, prediction_length)."""
    outputs_block = payload.get("outputs") or payload.get("forecasts") or []
    if not outputs_block:
        raise ValueError(f"Unrecognized TSFM API response: {payload}")

    first = outputs_block[0]
    rows: list[np.ndarray] = [ _flatten_forecast_values(first.get("mean"), prediction_length)]

    quantile_map: dict[float, np.ndarray] = {}
    for item in first.get("quantile_predictions", []):
        level = float(item.get("level", item.get("quantile", 0.5)))
        quantile_map[level] = _flatten_forecast_values(item.get("values"), prediction_length)

    for q in quantile_levels:
        if q in quantile_map:
            rows.append(quantile_map[q])
        else:
            nearest = min(quantile_map.keys(), key=lambda k: abs(k - q), default=None)
            if nearest is None:
                rows.append(rows[0].copy())
            else:
                rows.append(quantile_map[nearest])

    return np.stack(rows, axis=0).astype(np.float64)


def _flatten_forecast_values(values: Any, prediction_length: int) -> np.ndarray:
    """Convert nested API values like [[1.2], [1.3], ...] to 1D array."""
    if values is None:
        raise ValueError("Missing forecast values in TSFM API response")

    arr = np.asarray(values, dtype=np.float64)
    if arr.ndim == 2 and arr.shape[1] == 1:
        arr = arr[:, 0]
    elif arr.ndim > 1:
        arr = arr.reshape(arr.shape[0], -1).mean(axis=1)

    arr = arr.reshape(-1)
    if arr.size < prediction_length:
        raise ValueError(f"Forecast length {arr.size} < expected {prediction_length}")
    return arr[:prediction_length]