from typing import Iterable
import polars as pl


# Helpers
def _build_trading_calendar(
    dfs: Iterable[pl.DataFrame],
    symbol: str,
) -> pl.DataFrame:
    """
    Builds a master trading calendar using the union of all dates for a given symbol.
    """
    dates = (
        pl.concat(
            [
                df.select(["symbol", "date"])
                for df in dfs
                if "date" in df.columns
            ]
        )
        .filter(pl.col("symbol") == symbol)
        .unique()
        .sort("date")
    )

    return dates


def _forward_fill(
    df: pl.DataFrame,
    columns: list[str],
) -> pl.DataFrame:
    """
    Forward-fill selected columns in-place.
    """
    return df.with_columns(
        [pl.col(c).forward_fill() for c in columns]
    )


# Public API
def align_underlying_daily(
    *,
    symbol: str,
    ohlcv: pl.DataFrame,
    oi: pl.DataFrame,
    iv30: pl.DataFrame,
    iv60: pl.DataFrame,
    sofr: pl.DataFrame,
    usgg: pl.DataFrame,
    vix: pl.DataFrame,
) -> pl.DataFrame:
    """
    Align all normalized daily datasets into a single canonical
    underlying_daily table.

    Expected inputs:
    - All dataframes are already normalized
    - All contain (symbol, date)
    """

    # 1. Build master trading calendar
    calendar = _build_trading_calendar(
        dfs=[ohlcv, oi, iv30, iv60, sofr, usgg, vix],
        symbol=symbol,
    )

    # 2. Left-join everything onto calendar
    df = (
        calendar
        .join(ohlcv, on=["symbol", "date"], how="left")
        .join(oi, on=["symbol", "date"], how="left")
        .join(iv30, on=["symbol", "date"], how="left")
        .join(iv60, on=["symbol", "date"], how="left")
        .join(sofr, on=["symbol", "date"], how="left")
        .join(usgg, on=["symbol", "date"], how="left")
        .join(vix, on=["symbol", "date"], how="left")
        .sort("date")
    )

    # 3. Forward-fill macro variables
    df = _forward_fill(
        df,
        columns=[
            "rate_sofr_90d",
            "rate_usgg3m",
            "vix",
        ],
    )

    return df
