from __future__ import annotations
import polars as pl
import numpy as np


# Helpers
def _days_between(start: pl.Expr, end: pl.Expr) -> pl.Expr:
    """
    Calendar-day difference between two dates.
    """
    return (end - start).dt.total_days().cast(pl.Int64)

# Option chain enrichment
def enrich_option_chain_snapshot(
    *,
    chain: pl.DataFrame,
    underlying_row: dict,
) -> pl.DataFrame:
    """
    Enrich an option chain snapshot with derived columns.

    Adds:
    - time_to_expiry (days)
    - moneyness (K / S)
    - log_moneyness (log(K / S))
    """

    if "px_close" not in underlying_row:
        raise ValueError("underlying_row must contain px_close")

    spot = underlying_row["px_close"]

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

    return (
        chain
        .with_columns([
            # calendar days to expiry
            _days_between(pl.col("date"), pl.col("expiration"))
            .alias("time_to_expiry"),

            # moneyness
            (pl.col("strike") / spot)
            .cast(pl.Float64)
            .alias("moneyness"),

            # log-moneyness
            pl.when(pl.col("strike") > 0)
            .then(np.log(pl.col("strike") / spot))
            .otherwise(None)
            .cast(pl.Float64)
            .alias("log_moneyness"),
        ])
    )