from __future__ import annotations

import polars as pl
from typing import Dict

from ocf.data.normalize import (
    normalize_ohlcv,
    normalize_oi,
    normalize_ivol_30,
    normalize_ivol_60,
    normalize_sofr,
    normalize_usgg,
    normalize_vix,
    normalize_chain,
)

from ocf.data.align import align_underlying_daily
from ocf.data.enrich import enrich_option_chain_snapshot
from ocf.features.build import build_features
from ocf.greeks.api import per_option_greeks



# Normalize raw Bloomberg extracts
def step_normalize_raw(
    *,
    symbol: str,
    raw: Dict[str, str | pl.DataFrame],
    snapshot_date,
) -> Dict[str, pl.DataFrame]:
    """
    Normalizes raw Bloomberg inputs into canonical tables.
    """

    return {
        "ohlcv": normalize_ohlcv(raw["ohlcv"], symbol),
        "oi": normalize_oi(raw["oi"], symbol),
        "iv30": normalize_ivol_30(raw["iv30"], symbol),
        "iv60": normalize_ivol_60(raw["iv60"], symbol),
        "sofr": normalize_sofr(raw["sofr"], symbol),
        "usgg": normalize_usgg(raw["usgg"], symbol),
        "vix": normalize_vix(raw["vix"], symbol),
        "chain": normalize_chain(
            raw["chain"],
            symbol=symbol,
            snapshot_date=snapshot_date,
        ),
    }


# Align daily underlying data
def step_align_underlying(
    *,
    normalized: Dict[str, pl.DataFrame],
) -> pl.DataFrame:
    """
    Builds canonical_underlying_daily.
    """

    return align_underlying_daily(
        symbol=normalized["ohlcv"]["symbol"][0],
        ohlcv=normalized["ohlcv"],
        oi=normalized["oi"],
        iv30=normalized["iv30"],
        iv60=normalized["iv60"],
        sofr=normalized["sofr"],
        usgg=normalized["usgg"],
        vix=normalized["vix"],
    )


# Enrich option chain snapshot
def step_enrich_chain(
    *,
    chain: pl.DataFrame,
    underlying_daily: pl.DataFrame,
) -> pl.DataFrame:
    """
    Adds time_to_expiry, moneyness, log_moneyness.
    """
    row = underlying_daily.sort("date").tail(1).to_dicts()[0]

    return enrich_option_chain_snapshot(
        chain=chain,
        underlying_row=row,
    )


# Per-option Greeks
def step_per_option_greeks(
    *,
    chain_enriched: pl.DataFrame,
    underlying_daily: pl.DataFrame,
    rate_field: str = "rate_sofr_90d",
) -> pl.DataFrame:
    """
    Computes per-option Black–Scholes Greeks.
    """
    row = underlying_daily.sort("date").tail(1).to_dicts()[0]

    return per_option_greeks(
        chain_snapshot=chain_enriched,
        underlying_row=row,
        rate_field=rate_field,
    )


# Underlying feature engineering
def step_features_underlying(
    *,
    underlying_daily: pl.DataFrame,
    include=None,
    exclude=None,
) -> pl.DataFrame:
    """
    Builds model-ready underlying feature table.
    """

    return build_features(
        underlying_daily,
        include=include,
        exclude=exclude,
    )


# Final output bundle
def assemble_pipeline_outputs(
    *,
    underlying_daily: pl.DataFrame,
    features: pl.DataFrame,
    chain_snapshot: pl.DataFrame,
    per_option_greeks: pl.DataFrame,
) -> Dict[str, pl.DataFrame]:
    """
    Standard output contract for the pipeline.
    """

    return {
        "underlying_daily": underlying_daily,
        "features": features,
        "option_chain_snapshot": chain_snapshot,
        "per_option_greeks": per_option_greeks,
    }