from __future__ import annotations
import polars as pl
from typing import Iterable



# Helpers
def _assert_columns(df: pl.DataFrame, cols: Iterable[str], name: str):
    missing = set(cols) - set(df.columns)
    if missing:
        raise ValueError(f"{name} missing columns: {missing}")



# Underlying-level joins
def join_underlying_daily(
    *,
    ohlcv: pl.DataFrame,
    oi: pl.DataFrame | None = None,
    iv30: pl.DataFrame | None = None,
    iv60: pl.DataFrame | None = None,
    sofr: pl.DataFrame | None = None,
    usgg: pl.DataFrame | None = None,
    vix: pl.DataFrame | None = None,
) -> pl.DataFrame:
    """
    Join all underlying-level daily datasets.

    Joins:
    - OHLCV defines the trading calendar
    - everything else left-joins onto OHLCV
    """

    _assert_columns(ohlcv, ["symbol", "date"], "ohlcv")

    df = ohlcv.clone()

    def _left_join(
        left: pl.DataFrame,
        right: pl.DataFrame | None,
        name: str,
    ) -> pl.DataFrame:
        if right is None:
            return left
        _assert_columns(right, ["symbol", "date"], name)
        return left.join(
            right,
            on=["symbol", "date"],
            how="left",
            suffix=f"__{name}",
        )

    df = _left_join(df, oi, "oi")
    df = _left_join(df, iv30, "iv30")
    df = _left_join(df, iv60, "iv60")
    df = _left_join(df, sofr, "sofr")
    df = _left_join(df, usgg, "usgg")
    df = _left_join(df, vix, "vix")

    return df



# Option-level joins
def join_options_with_greeks(
    *,
    chain: pl.DataFrame,
    greeks: pl.DataFrame,
) -> pl.DataFrame:
    """
    Join normalized chain metadata with per-option Greeks.

    Joins:
    - chain provides metadata only
    - inner join to avoid phantom options
    """

    _assert_columns(
        chain,
        ["symbol", "date", "expiration", "strike", "call_ticker", "put_ticker"],
        "chain",
    )

    _assert_columns(
        greeks,
        ["symbol", "date", "expiration", "strike", "option_type"],
        "greeks",
    )

    # Normalize ticker column for join
    greeks = greeks.with_columns(
        pl.when(pl.col("option_type") == "call")
        .then(pl.col("ticker"))
        .otherwise(pl.col("ticker"))
        .alias("ticker")
    )

    chain_long = (
        chain
        .select([
            "symbol",
            "date",
            "expiration",
            "strike",
            pl.col("call_ticker").alias("ticker"),
        ])
        .with_columns(pl.lit("call").alias("option_type"))
        .vstack(
            chain.select([
                "symbol",
                "date",
                "expiration",
                "strike",
                pl.col("put_ticker").alias("ticker"),
            ])
            .with_columns(pl.lit("put").alias("option_type"))
        )
    )

    return greeks.join(
        chain_long,
        on=[
            "symbol",
            "date",
            "expiration",
            "strike",
            "ticker",
            "option_type",
        ],
        how="inner",
    )