from __future__ import annotations
from typing import Dict

import polars as pl

from ocf.greeks.per_option import build_per_option_greeks
from ocf.greeks.exposure_curves import build_exposure_features


# Per-option Greeks
def per_option_greeks(
    *,
    chain_snapshot: pl.DataFrame,
    underlying_row: dict | pl.DataFrame,
    rate_field: str = "rate_sofr_90d",
) -> pl.DataFrame:
    """
    Computes Black–Scholes Greeks per option contract.
    """
    greeks = build_per_option_greeks(
        chain=chain_snapshot,
        underlying_row=underlying_row,
        rate_field=rate_field,
    )

    if greeks.is_empty():
        return greeks

    for col in ("delta", "gamma", "vega", "theta"):
        if col not in greeks.columns:
            raise RuntimeError(
                f"Per-option Greeks missing column: {col}"
            )

    return greeks


# Aggregate exposure features (ML-ready)
def exposure_features(
    greeks_df: pl.DataFrame,
    *,
    atm_tau: float = 0.15,
    slope_window: float = 0.25,
) -> pl.DataFrame:
    """
    Computes aggregate exposure features from per-option Greeks.
    """
    if greeks_df.is_empty():
        return pl.DataFrame([])

    features = build_exposure_features(
        greeks_df,
        atm_tau=atm_tau,
        slope_window=slope_window,
    )

    return pl.DataFrame([features])


# Snapshot-level convenience API
def option_chain_greeks(
    *,
    chain_snapshot: pl.DataFrame,
    underlying_row: dict | pl.DataFrame,
    rate_field: str = "rate_sofr_90d",
    return_per_option: bool = True,
    return_exposures: bool = True,
    atm_tau: float = 0.15,
    slope_window: float = 0.25,
) -> Dict[str, pl.DataFrame]:
    """
    Computes a full Greeks view for a single option chain snapshot.
    """
    outputs: Dict[str, pl.DataFrame] = {}

    greeks_df = per_option_greeks(
        chain_snapshot=chain_snapshot,
        underlying_row=underlying_row,
        rate_field=rate_field,
    )

    if return_per_option:
        outputs["per_option_greeks"] = greeks_df

    if return_exposures:
        outputs["exposure_features"] = exposure_features(
            greeks_df,
            atm_tau=atm_tau,
            slope_window=slope_window,
        )

    return outputs