from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from typing import Any

import polars as pl

from uforia.config import AppConfig, get_config
from uforia.storage import DatasetName, DatasetStore


@dataclass(frozen=True)
class RazorLegSnapshot:
    ts: datetime
    underlying: str
    expiry: datetime
    call_name: str
    put_name: str
    strike: float
    underlying_price: float
    straddle_value: float
    straddle_delta: float
    atm_iv_7d: float | None
    quality_flag: str


def compute_razor_pnl(
    store: DatasetStore,
    *,
    start: datetime,
    end: datetime,
    underlying: str,
    tenor_days: int = 7,
) -> pl.DataFrame:
    snapshots = _build_reference_straddles(
        store,
        start=start,
        end=end,
        underlying=underlying,
        tenor_days=tenor_days,
    )
    if snapshots.is_empty():
        return pl.DataFrame(
            schema={
                "ts": pl.Datetime(time_zone="UTC"),
                "date": pl.Date,
                "underlying": pl.String,
                "reference_tenor_days": pl.Int64,
                "quality_flag": pl.String,
                "buyer_increment": pl.Float64,
                "seller_increment": pl.Float64,
                "option_reval_pnl": pl.Float64,
                "hedge_pnl": pl.Float64,
                "roll_pnl": pl.Float64,
                "buyer_cumulative": pl.Float64,
                "seller_cumulative": pl.Float64,
                "razor_pnl": pl.Float64,
            }
        )

    options = store.read(
        DatasetName.OPTION_CHAIN_SNAPSHOT,
        start=start,
        end=end,
        filters={"underlying": underlying},
    )
    option_map: dict[tuple[datetime, str], float] = {
        (row["ts"], row["instrument_name"]): float(row["mid"])
        for row in options.select(["ts", "instrument_name", "mid"])
        .drop_nulls(["mid"])
        .iter_rows(named=True)
    }

    rows = snapshots.sort("ts").to_dicts()
    buyer_cumulative = 0.0
    seller_cumulative = 0.0
    results: list[dict[str, Any]] = []
    for idx, row in enumerate(rows):
        buyer_increment: float | None = None
        seller_increment: float | None = None
        option_reval_pnl: float | None = None
        hedge_pnl: float | None = None
        roll_pnl: float | None = None
        if idx > 0:
            prev = rows[idx - 1]
            current_call = option_map.get((row["ts"], prev["call_name"]))
            current_put = option_map.get((row["ts"], prev["put_name"]))
            if current_call is not None and current_put is not None:
                current_previous_value = current_call + current_put
                option_reval_pnl = current_previous_value - float(prev["straddle_value"])
                hedge_pnl = -float(prev["straddle_delta"]) * (
                    float(row["underlying_price"]) - float(prev["underlying_price"])
                )
                if prev["call_name"] != row["call_name"] or prev["put_name"] != row["put_name"]:
                    roll_pnl = -(float(row["straddle_value"]) - current_previous_value)
                else:
                    roll_pnl = 0.0
                buyer_increment = option_reval_pnl + hedge_pnl + roll_pnl
                seller_increment = -buyer_increment
                buyer_cumulative += buyer_increment
                seller_cumulative += seller_increment
        quality_flag = str(row["quality_flag"])
        if buyer_increment is None:
            quality_flag = "initializing"
        results.append(
            {
                "ts": row["ts"],
                "date": row["date"],
                "underlying": row["underlying"],
                "reference_tenor_days": row["reference_tenor_days"],
                "quality_flag": quality_flag,
                "buyer_increment": buyer_increment,
                "seller_increment": seller_increment,
                "option_reval_pnl": option_reval_pnl,
                "hedge_pnl": hedge_pnl,
                "roll_pnl": roll_pnl,
                "buyer_cumulative": buyer_cumulative if buyer_increment is not None else None,
                "seller_cumulative": seller_cumulative if seller_increment is not None else None,
                "razor_pnl": (seller_cumulative - buyer_cumulative)
                if buyer_increment is not None
                else None,
            }
        )
    return pl.DataFrame(results)


def compute_razor_vrp(
    store: DatasetStore,
    *,
    start: datetime,
    end: datetime,
    underlying: str,
    tenor_days: int = 7,
) -> pl.DataFrame:
    surface = store.read(
        DatasetName.SURFACE_FEATURES,
        start=start,
        end=end,
        filters={"underlying": underlying, "source": "ssvi"},
    ).sort("ts")
    options = store.read(
        DatasetName.OPTION_CHAIN_SNAPSHOT,
        start=start,
        end=end,
        filters={"underlying": underlying},
    )
    if surface.is_empty() or options.is_empty():
        return pl.DataFrame()

    spot = (
        options.group_by("ts")
        .agg(pl.col("underlying_price").median().alias("underlying_price"))
        .sort("ts")
        .with_columns(
            pl.col("underlying_price").log().diff().alias("log_return"),
            (pl.col("underlying_price").log().diff() ** 2).alias("squared_return"),
        )
    )
    periods = tenor_days * 24
    realized = spot.with_columns(
        (
            pl.col("squared_return")
            .rolling_mean(window_size=periods, min_samples=max(2, periods // 2))
            * 24.0
            * 365.0
        ).alias("realized_var_7d")
    ).select(["ts", "realized_var_7d"])
    joined = surface.join(realized, on="ts", how="left").with_columns(
        (pl.col("atm_iv_7d") ** 2).alias("implied_var_7d"),
        (((pl.col("atm_iv_7d") ** 2) - pl.col("realized_var_7d")) / (24.0 * 365.0) * 10_000.0)
        .alias("vrp_increment"),
    )
    cumulative = joined.with_columns(
        pl.col("vrp_increment").fill_null(0.0).cum_sum().alias("razor_vrp")
    )
    return cumulative.select(
        [
            "ts",
            "date",
            "underlying",
            pl.lit(tenor_days).alias("reference_tenor_days"),
            pl.col("implied_var_7d"),
            pl.col("realized_var_7d"),
            pl.col("vrp_increment"),
            "razor_vrp",
            "quality_flag",
        ]
    )


def build_razor_series(
    store: DatasetStore,
    *,
    start: datetime,
    end: datetime,
    underlying: str,
    config: AppConfig | None = None,
) -> pl.DataFrame:
    cfg = config or get_config()
    tenor_days = cfg.modeling.razor_reference_tenor_days
    pnl = compute_razor_pnl(
        store,
        start=start,
        end=end,
        underlying=underlying,
        tenor_days=tenor_days,
    )
    vrp = compute_razor_vrp(
        store,
        start=start,
        end=end,
        underlying=underlying,
        tenor_days=tenor_days,
    )
    if pnl.is_empty() and vrp.is_empty():
        return pl.DataFrame()
    if pnl.is_empty():
        frame = vrp.with_columns(
            pl.lit(None, dtype=pl.Float64).alias("razor_pnl"),
            pl.lit(None, dtype=pl.Float64).alias("buyer_increment"),
            pl.lit(None, dtype=pl.Float64).alias("seller_increment"),
            pl.lit(None, dtype=pl.Float64).alias("option_reval_pnl"),
            pl.lit(None, dtype=pl.Float64).alias("hedge_pnl"),
            pl.lit(None, dtype=pl.Float64).alias("roll_pnl"),
        )
    elif vrp.is_empty():
        frame = pnl.with_columns(pl.lit(None, dtype=pl.Float64).alias("razor_vrp"))
    else:
        frame = pnl.join(
            vrp.select(["ts", "razor_vrp"]),
            on="ts",
            how="full",
            coalesce=True,
        ).sort("ts")
    if "date" not in frame.columns:
        frame = frame.with_columns(pl.col("ts").dt.date().alias("date"))
    if "underlying" not in frame.columns:
        frame = frame.with_columns(pl.lit(underlying).alias("underlying"))
    if "reference_tenor_days" not in frame.columns:
        frame = frame.with_columns(pl.lit(tenor_days).alias("reference_tenor_days"))
    if "quality_flag" not in frame.columns:
        frame = frame.with_columns(pl.lit("ok").alias("quality_flag"))

    lookback = cfg.modeling.razor_zscore_lookback_days * 24
    frame = frame.with_columns(
        _rolling_zscore("razor_pnl", lookback).alias("razor_pnl_zscore"),
        _rolling_zscore("razor_vrp", lookback).alias("razor_vrp_zscore"),
    ).with_columns(
        pl.lit(cfg.modeling.razor_source).alias("source"),
        (pl.col("razor_pnl_zscore") >= cfg.modeling.razor_long_vol_threshold).alias(
            "long_vol_signal"
        ),
        (pl.col("razor_pnl_zscore") <= cfg.modeling.razor_short_vol_threshold).alias(
            "short_vol_signal"
        ),
    ).with_columns(
        pl.when(pl.col("long_vol_signal"))
        .then(pl.lit("double_buy"))
        .when(pl.col("short_vol_signal"))
        .then(pl.lit("double_sell"))
        .otherwise(pl.lit("neutral"))
        .alias("trade_band")
    ).with_columns(
        pl.when(pl.col("razor_pnl").is_null() | pl.col("razor_pnl_zscore").is_null())
        .then(pl.lit("initializing"))
        .otherwise(pl.lit("ok"))
        .alias("quality_flag")
    )
    return frame.select(
        [
            "ts",
            "date",
            "underlying",
            "source",
            "reference_tenor_days",
            "razor_pnl",
            "razor_vrp",
            "buyer_increment",
            "seller_increment",
            "option_reval_pnl",
            "hedge_pnl",
            "roll_pnl",
            "razor_pnl_zscore",
            "razor_vrp_zscore",
            "trade_band",
            "long_vol_signal",
            "short_vol_signal",
            "quality_flag",
        ]
    )


def _build_reference_straddles(
    store: DatasetStore,
    *,
    start: datetime,
    end: datetime,
    underlying: str,
    tenor_days: int,
) -> pl.DataFrame:
    options = store.read(
        DatasetName.OPTION_CHAIN_SNAPSHOT,
        start=start,
        end=end,
        filters={"underlying": underlying},
    )
    if options.is_empty():
        return pl.DataFrame()

    records: list[dict[str, Any]] = []
    for ts_key, snapshot in options.partition_by("ts", maintain_order=True, as_dict=True).items():
        ts = ts_key[0] if isinstance(ts_key, tuple) else ts_key
        snapshot = snapshot.drop_nulls(["mid", "strike", "expiry", "delta", "underlying_price"])
        if snapshot.is_empty():
            continue
        expiries = (
            snapshot.with_columns(
                ((pl.col("expiry") - pl.lit(ts)).dt.total_seconds() / 86400.0).alias(
                    "days_to_expiry"
                )
            )
            .filter(pl.col("days_to_expiry") > 0)
            .group_by("expiry")
            .agg(pl.first("days_to_expiry").alias("days_to_expiry"))
            .with_columns((pl.col("days_to_expiry") - tenor_days).abs().alias("distance"))
            .sort(["distance", "days_to_expiry"])
        )
        if expiries.is_empty():
            continue
        expiry = expiries.item(0, "expiry")
        tenor_snapshot = snapshot.filter(pl.col("expiry") == expiry)
        forward = float(
            tenor_snapshot.select(pl.col("forward").drop_nulls().median()).item(0, 0)
            or tenor_snapshot.select(pl.col("underlying_price").median()).item(0, 0)
        )
        pairs = (
            tenor_snapshot.group_by("strike")
            .agg(
                pl.when(pl.col("option_type") == "call")
                .then(pl.col("mid"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("call_mid"),
                pl.when(pl.col("option_type") == "put")
                .then(pl.col("mid"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("put_mid"),
                pl.when(pl.col("option_type") == "call")
                .then(pl.col("instrument_name"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("call_name"),
                pl.when(pl.col("option_type") == "put")
                .then(pl.col("instrument_name"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("put_name"),
                pl.when(pl.col("option_type") == "call")
                .then(pl.col("delta"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("call_delta"),
                pl.when(pl.col("option_type") == "put")
                .then(pl.col("delta"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("put_delta"),
                pl.when(pl.col("option_type") == "call")
                .then(pl.col("mark_iv"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("call_iv"),
                pl.when(pl.col("option_type") == "put")
                .then(pl.col("mark_iv"))
                .otherwise(None)
                .drop_nulls()
                .first()
                .alias("put_iv"),
            )
            .drop_nulls(["call_mid", "put_mid", "call_name", "put_name"])
            .with_columns((pl.col("strike") - forward).abs().alias("distance"))
            .sort("distance")
        )
        if pairs.is_empty():
            continue
        pair = pairs.row(0, named=True)
        atm_iv_7d = None
        if pair["call_iv"] is not None and pair["put_iv"] is not None:
            atm_iv_7d = (float(pair["call_iv"]) + float(pair["put_iv"])) / 2.0
        records.append(
            {
                "ts": ts,
                "date": ts.date(),
                "underlying": underlying,
                "reference_tenor_days": tenor_days,
                "expiry": expiry,
                "call_name": str(pair["call_name"]),
                "put_name": str(pair["put_name"]),
                "strike": float(pair["strike"]),
                "underlying_price": float(tenor_snapshot.item(0, "underlying_price")),
                "straddle_value": float(pair["call_mid"]) + float(pair["put_mid"]),
                "straddle_delta": float(pair["call_delta"] or 0.0)
                + float(pair["put_delta"] or 0.0),
                "atm_iv_7d": atm_iv_7d,
                "quality_flag": "ok",
            }
        )
    if not records:
        return pl.DataFrame(
            schema={
                "ts": pl.Datetime(time_zone="UTC"),
                "date": pl.Date,
                "underlying": pl.String,
                "reference_tenor_days": pl.Int64,
                "expiry": pl.Datetime(time_zone="UTC"),
                "call_name": pl.String,
                "put_name": pl.String,
                "strike": pl.Float64,
                "underlying_price": pl.Float64,
                "straddle_value": pl.Float64,
                "straddle_delta": pl.Float64,
                "atm_iv_7d": pl.Float64,
                "quality_flag": pl.String,
            }
        )
    return pl.DataFrame(records).sort("ts")


def _rolling_zscore(column: str, window_size: int) -> pl.Expr:
    min_samples = max(4, window_size // 3)
    mean = pl.col(column).rolling_mean(window_size=window_size, min_samples=min_samples)
    std = pl.col(column).rolling_std(window_size=window_size, min_samples=min_samples)
    return ((pl.col(column) - mean) / std).fill_nan(None)
