from __future__ import annotations

import numpy as np
import polars as pl
from datetime import date

from ocf.greeks.vol_mapping import map_sigma
from ocf.greeks.black_scholes import bs_greeks

# Validation
REQUIRED_CHAIN_COLS = {
    "symbol",
    "date",
    "expiration",
    "strike",
    "call_ticker",
    "put_ticker",
    "contract_size",
}

def _validate_chain(df: pl.DataFrame):
    missing = REQUIRED_CHAIN_COLS - set(df.columns)
    if missing:
        raise ValueError(f"Missing required chain columns: {missing}")

# Helpers
def _days_to_expiry(valuation_date: date, expiry: date) -> int:
    return max((expiry - valuation_date).days, 0)

# Core builder
def build_per_option_greeks(
    *,
    chain: pl.DataFrame,
    underlying_row: dict,
    rate_field: str = "rate_sofr_90d",
) -> pl.DataFrame:
    """
    Builds per-option Greeks for a chain snapshot.
    """

    _validate_chain(chain)

    S = underlying_row["px_close"]
    r = underlying_row.get(rate_field, np.nan)
    valuation_date = underlying_row["date"]

    if S <= 0:
        raise ValueError("Underlying spot price must be positive")

    rows = []

    for opt in chain.iter_rows(named=True):
        expiry = opt["expiration"]
        K = opt["strike"]
        contract_size = opt["contract_size"]

        T_days = _days_to_expiry(valuation_date, expiry)
        if T_days <= 0:
            continue

        T = T_days / 365.0

        sigma = map_sigma(
            row=underlying_row,
            strike=K,
            spot=S,
            days_to_expiry=T_days,
        )

        # CALL
        g = bs_greeks(
            S=S,
            K=K,
            T=T,
            r=r,
            sigma=sigma,
            is_call=True,
        )

        rows.append({
            "symbol": opt["symbol"],
            "date": valuation_date,
            "expiration": expiry,
            "option_type": "call",
            "ticker": opt["call_ticker"],
            "strike": K,
            "days_to_expiry": T_days,
            "sigma": sigma,
            "price": g.price,
            "delta": g.delta,
            "gamma": g.gamma,
            "vega": g.vega,
            "theta": g.theta,
            "charm": g.charm,
            "vanna": g.vanna,
            "vomma": g.vomma,
            "contract_size": contract_size,
        })

        # PUT
        g = bs_greeks(
            S=S,
            K=K,
            T=T,
            r=r,
            sigma=sigma,
            is_call=False,
        )

        rows.append({
            "symbol": opt["symbol"],
            "date": valuation_date,
            "expiration": expiry,
            "option_type": "put",
            "ticker": opt["put_ticker"],
            "strike": K,
            "days_to_expiry": T_days,
            "sigma": sigma,
            "price": g.price,
            "delta": g.delta,
            "gamma": g.gamma,
            "vega": g.vega,
            "theta": g.theta,
            "charm": g.charm,
            "vanna": g.vanna,
            "vomma": g.vomma,
            "contract_size": contract_size,
        })

    if not rows:
        return pl.DataFrame([])

    return pl.DataFrame(rows)