from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Protocol

import numpy as np
import polars as pl

from uforia.storage import DatasetName, DatasetStore


class VolIndexProvider(Protocol):
    source: str

    def load_series(self, store: DatasetStore, start: datetime, end: datetime) -> pl.DataFrame: ...

    def build_point(self, store: DatasetStore, ts: datetime) -> dict[str, object]: ...


@dataclass
class DVOLProvider:
    underlying: str = "BTC"
    source: str = "dvol"

    def load_series(self, store: DatasetStore, start: datetime, end: datetime) -> pl.DataFrame:
        frame = store.read(
            DatasetName.DVOL_SNAPSHOT,
            start=start,
            end=end,
            filters={"underlying": self.underlying, "source": self.source},
        )
        if frame.is_empty():
            return pl.DataFrame(
                schema={
                    "ts": pl.Datetime(time_zone="UTC"),
                    "date": pl.Date,
                    "underlying": pl.String,
                    "source": pl.String,
                    "value": pl.Float64,
                    "quality_flag": pl.String,
                }
            )
        return frame.select(
            [
                "ts",
                "date",
                "underlying",
                pl.lit(self.source).alias("source"),
                pl.col("dvol_close").alias("value"),
                "quality_flag",
            ]
        )

    def build_point(self, store: DatasetStore, ts: datetime) -> dict[str, object]:
        frame = self.load_series(store, ts - timedelta(hours=1), ts + timedelta(hours=1))
        if frame.is_empty():
            raise ValueError(f"No DVOL data found around {ts.isoformat()}")
        row = frame.sort("ts").tail(1).row(0, named=True)
        return dict(row)


@dataclass
class BVIXProvider:
    underlying: str = "BTC"
    source: str = "bvix"

    def load_series(self, store: DatasetStore, start: datetime, end: datetime) -> pl.DataFrame:
        options = store.read(
            DatasetName.OPTION_CHAIN_SNAPSHOT,
            start=start,
            end=end,
            filters={"underlying": self.underlying},
        )
        if options.is_empty():
            return pl.DataFrame(
                schema={
                    "ts": pl.Datetime(time_zone="UTC"),
                    "date": pl.Date,
                    "underlying": pl.String,
                    "source": pl.String,
                    "value": pl.Float64,
                    "quality_flag": pl.String,
                }
            )

        rows: list[dict[str, object]] = []
        for ts, snapshot in options.partition_by("ts", maintain_order=True, as_dict=True).items():
            try:
                value = _compute_bvix(snapshot)
                quality_flag = "ok"
            except ValueError:
                value = None
                quality_flag = "insufficient_data"
            rows.append(
                {
                    "ts": ts[0] if isinstance(ts, tuple) else ts,
                    "date": (ts[0] if isinstance(ts, tuple) else ts).date(),
                    "underlying": self.underlying,
                    "source": self.source,
                    "value": value,
                    "quality_flag": quality_flag,
                }
            )
        return pl.DataFrame(rows).sort("ts")

    def build_point(self, store: DatasetStore, ts: datetime) -> dict[str, object]:
        frame = self.load_series(store, ts - timedelta(hours=1), ts + timedelta(hours=1))
        if frame.is_empty():
            raise ValueError(f"No option-chain data found around {ts.isoformat()}")
        row = frame.sort("ts").tail(1).row(0, named=True)
        return dict(row)


def build_vol_index(
    store: DatasetStore,
    start: datetime,
    end: datetime,
    source: str,
    underlying: str = "BTC",
) -> pl.DataFrame:
    provider = get_provider(source, underlying=underlying)
    return provider.load_series(store, start, end)


def get_provider(source: str, underlying: str = "BTC") -> VolIndexProvider:
    if source == "dvol":
        return DVOLProvider(underlying=underlying)
    if source == "bvix":
        return BVIXProvider(underlying=underlying)
    raise ValueError(f"Unsupported vol index source: {source}")


def _compute_bvix(snapshot: pl.DataFrame) -> float:
    ts = snapshot.item(0, "ts")
    target_days = 30
    expiries = (
        snapshot.select(
            [
                "expiry",
                ((pl.col("expiry") - pl.lit(ts)).dt.total_seconds() / 86400.0).alias(
                    "days_to_expiry"
                ),
            ]
        )
        .unique("expiry")
        .sort("days_to_expiry")
    )
    near = expiries.filter(pl.col("days_to_expiry") <= target_days).tail(1)
    nxt = expiries.filter(pl.col("days_to_expiry") >= target_days).head(1)
    if near.is_empty() or nxt.is_empty():
        raise ValueError("Could not bracket 30-day expiry")

    near_expiry = near.item(0, "expiry")
    next_expiry = nxt.item(0, "expiry")
    near_var = _variance_for_expiry(snapshot.filter(pl.col("expiry") == near_expiry), ts)
    next_var = _variance_for_expiry(snapshot.filter(pl.col("expiry") == next_expiry), ts)
    t1 = near.item(0, "days_to_expiry") / 365.0
    t2 = nxt.item(0, "days_to_expiry") / 365.0
    target = target_days / 365.0
    interpolated = _interpolate_constant_maturity_variance(near_var, next_var, t1, t2, target)
    if interpolated < 0:
        raise ValueError("Interpolated variance is negative")
    return 100.0 * float(np.sqrt(interpolated))


def _interpolate_constant_maturity_variance(
    near_var: float,
    next_var: float,
    near_t: float,
    next_t: float,
    target_t: float,
) -> float:
    if target_t <= 0:
        raise ValueError("Target maturity must be positive")
    if abs(next_t - near_t) < 1e-12:
        return near_var
    weighted = (
        near_t * near_var * (next_t - target_t) + next_t * next_var * (target_t - near_t)
    ) / (next_t - near_t)
    return weighted / target_t


def _variance_for_expiry(expiry_frame: pl.DataFrame, ts: datetime) -> float:
    t = max(
        (expiry_frame.item(0, "expiry") - ts).total_seconds() / (365.0 * 24 * 3600),
        1 / (365.0 * 24),
    )
    strikes = sorted(expiry_frame["strike"].unique().drop_nulls().to_list())
    if len(strikes) < 3:
        raise ValueError("Not enough strikes")

    parity_input = expiry_frame.with_columns(
        (pl.col("mid") * pl.col("index_price")).alias("mid_usd")
    )
    parity = (
        parity_input.pivot(
            index="strike", on="option_type", values="mid_usd", aggregate_function="first"
        )
        .drop_nulls(subset=["call", "put"])
        .with_columns((pl.col("call") - pl.col("put")).abs().alias("abs_diff"))
        .sort("abs_diff")
    )
    if parity.is_empty():
        raise ValueError("No valid call/put parity pairs")
    parity_row = parity.row(0, named=True)
    k_star = float(parity_row["strike"])
    rate = float(expiry_frame.item(0, "interest_rate") or 0.0)
    forward = k_star + np.exp(rate * t) * (float(parity_row["call"]) - float(parity_row["put"]))
    strikes_below = [strike for strike in strikes if strike <= forward]
    if not strikes_below:
        raise ValueError("No strike below forward")
    k0 = max(strikes_below)

    sorted_strikes = np.array(strikes, dtype=float)
    delta_k = np.zeros_like(sorted_strikes)
    delta_k[0] = sorted_strikes[1] - sorted_strikes[0]
    delta_k[-1] = sorted_strikes[-1] - sorted_strikes[-2]
    delta_k[1:-1] = (sorted_strikes[2:] - sorted_strikes[:-2]) / 2.0

    quote_map = {
        (float(row["strike"]), row["option_type"]): float(row["mid"]) * float(row["index_price"])
        for row in expiry_frame.select(["strike", "option_type", "mid", "index_price"])
        .drop_nulls()
        .iter_rows(named=True)
    }
    variance_sum = 0.0
    for strike, dk in zip(sorted_strikes, delta_k, strict=True):
        option_type = "put" if strike <= k0 else "call"
        premium = quote_map.get((float(strike), option_type))
        if premium is None or premium <= 0:
            continue
        variance_sum += (dk / (strike**2)) * np.exp(rate * t) * premium

    return (2.0 / t) * variance_sum - (1.0 / t) * ((forward / k0) - 1.0) ** 2
