from pathlib import Path
import polars as pl
from typing import Set


# Core helpers
def _read_csv(obj: str | Path | pl.DataFrame) -> pl.DataFrame:
    """
    Read CSV from path OR passthrough DataFrame.
    """
    if isinstance(obj, pl.DataFrame):
        return obj.clone()

    if isinstance(obj, (str, Path)):
        return pl.read_csv(obj, ignore_errors=False)

    raise TypeError(f"Unsupported input type for CSV reader: {type(obj)}")


def _add_symbol(df: pl.DataFrame, symbol: str) -> pl.DataFrame:
    if "symbol" in df.columns:
        return df
    return df.with_columns(pl.lit(symbol).alias("symbol"))


def _parse_bloomberg_date(expr: pl.Expr) -> pl.Expr:
    """
    Bloomberg dates: DD/MM/YY
    """
    return (
        expr
        .str.strptime(pl.Date, "%d/%m/%y", strict=False)
        .map_elements(
            lambda d: d.replace(year=d.year + 2000) if d.year < 1970 else d
        )
    )


def _has_columns(df: pl.DataFrame, cols: Set[str]) -> bool:
    return cols.issubset(set(df.columns))


# IVOL 30D (moneyness)
_IV30_COLS = {
    "date",
    "iv_30_m90", "iv_30_m95", "iv_30_m100", "iv_30_m105", "iv_30_m110",
}

def normalize_ivol_30(obj: str | Path | pl.DataFrame, symbol: str) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _IV30_COLS):
        return _add_symbol(df, symbol).sort("date")

    return (
        df.rename({
            "Date": "date",
            "IVOL_30_90": "iv_30_m90",
            "IVOL_30_95": "iv_30_m95",
            "IVOL_30_100": "iv_30_m100",
            "IVOL_30_105": "iv_30_m105",
            "IVOL_30_110": "iv_30_m110",
        })
        .with_columns(
            _parse_bloomberg_date(pl.col("date")).alias("date"),
            pl.col("^iv_30_.*$").cast(pl.Float64),
        )
        .pipe(_add_symbol, symbol)
        .sort("date")
    )


# IVOL 60D (moneyness)
_IV60_COLS = {
    "date",
    "iv_60_m90", "iv_60_m95", "iv_60_m100", "iv_60_m105", "iv_60_m110",
}

def normalize_ivol_60(obj: str | Path | pl.DataFrame, symbol: str) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _IV60_COLS):
        return _add_symbol(df, symbol).sort("date")

    return (
        df.rename({
            "Date": "date",
            "IVOL_60_90": "iv_60_m90",
            "IVOL_60_95": "iv_60_m95",
            "IVOL_60_100": "iv_60_m100",
            "IVOL_60_105": "iv_60_m105",
            "IVOL_60_110": "iv_60_m110",
        })
        .with_columns(
            _parse_bloomberg_date(pl.col("date")).alias("date"),
            pl.col("^iv_60_.*$").cast(pl.Float64),
        )
        .pipe(_add_symbol, symbol)
        .sort("date")
    )


# OHLCV
_OHLCV_COLS = {
    "date", "px_open", "px_high", "px_low", "px_close", "volume",
}

def normalize_ohlcv(obj: str | Path | pl.DataFrame, symbol: str) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _OHLCV_COLS):
        return _add_symbol(df, symbol).sort("date")

    return (
        df.rename({
            "Date": "date",
            "Open": "px_open",
            "High": "px_high",
            "Low": "px_low",
            "Close": "px_close",
            "Volume": "volume",
        })
        .with_columns(
            _parse_bloomberg_date(pl.col("date")).alias("date"),
            pl.col("px_open").cast(pl.Float64),
            pl.col("px_high").cast(pl.Float64),
            pl.col("px_low").cast(pl.Float64),
            pl.col("px_close").cast(pl.Float64),
            pl.col("volume").cast(pl.Float64),
        )
        .pipe(_add_symbol, symbol)
        .sort("date")
    )


# OI + ATM IV
_OI_COLS = {
    "date", "px_close", "bid", "ask", "mid",
    "iv_atm_hist", "call_open_int_tot", "put_open_int_tot",
}

def normalize_oi(obj: str | Path | pl.DataFrame, symbol: str) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _OI_COLS):
        return _add_symbol(df, symbol).sort("date")

    return (
        df.rename({
            "Date": "date",
            "Close": "px_close",
            "Bid": "bid",
            "Ask": "ask",
            "Total Call Open Interest": "call_open_int_tot",
            "Total Put Open Interest": "put_open_int_tot",
            "Historical Call Implied Volatility": "iv_call_hist",
            "Historical Put Implied Volatility": "iv_put_hist",
        })
        .with_columns(
            _parse_bloomberg_date(pl.col("date")).alias("date"),
            pl.col("bid").cast(pl.Float64),
            pl.col("ask").cast(pl.Float64),
            ((pl.col("bid") + pl.col("ask")) / 2).alias("mid"),
            pl.col("iv_call_hist").cast(pl.Float64).alias("iv_atm_hist"),
            pl.col("call_open_int_tot").cast(pl.Float64),
            pl.col("put_open_int_tot").cast(pl.Float64),
        )
        .drop(["iv_call_hist", "iv_put_hist"])
        .pipe(_add_symbol, symbol)
        .sort("date")
    )


# Risk-free SOFR
_SOFR_COLS = {"date", "rate_sofr_90d"}

def normalize_sofr(obj: str | Path | pl.DataFrame, symbol: str) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _SOFR_COLS):
        return _add_symbol(df, symbol).sort("date")

    return (
        df.rename({
            "Date": "date",
            "Risk Free SOFR": "rate_sofr_90d",
        })
        .with_columns(
            _parse_bloomberg_date(pl.col("date")).alias("date"),
            (pl.col("rate_sofr_90d") / 100).cast(pl.Float64),
        )
        .pipe(_add_symbol, symbol)
        .sort("date")
    )


# Risk-free USGG
_USGG_COLS = {"date", "rate_usgg3m"}

def normalize_usgg(obj: str | Path | pl.DataFrame, symbol: str) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _USGG_COLS):
        return _add_symbol(df, symbol).sort("date")

    return (
        df.rename({
            "Date": "date",
            "Risk Free USGG": "rate_usgg3m",
        })
        .with_columns(
            _parse_bloomberg_date(pl.col("date")).alias("date"),
            (pl.col("rate_usgg3m") / 100).cast(pl.Float64),
        )
        .pipe(_add_symbol, symbol)
        .sort("date")
    )


# VIX
_VIX_COLS = {"date", "vix"}

def normalize_vix(obj: str | Path | pl.DataFrame, symbol: str) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _VIX_COLS):
        return _add_symbol(df, symbol).sort("date")

    return (
        df.rename({
            "Date": "date",
            "Close": "vix",
        })
        .with_columns(
            _parse_bloomberg_date(pl.col("date")).alias("date"),
            pl.col("vix").cast(pl.Float64),
        )
        .pipe(_add_symbol, symbol)
        .sort("date")
    )


# Option chain snapshot
_CHAIN_COLS = {
    "symbol", "date", "call_ticker", "put_ticker",
    "expiration", "strike", "exercise_type", "contract_size", "periodicity",
}

def normalize_chain(
    obj: str | Path | pl.DataFrame,
    *,
    symbol: str,
    snapshot_date,
) -> pl.DataFrame:
    df = _read_csv(obj)

    if _has_columns(df, _CHAIN_COLS):
        return df.sort(["expiration", "strike"])

    return (
        df.rename({
            "Call Ticker": "call_ticker",
            "Put Ticker": "put_ticker",
            "Expiration": "expiration",
            "Strike": "strike",
            "Exercise Type": "exercise_type",
            "Contract Size": "contract_size",
            "Periodicity": "periodicity",
        })
        .with_columns(
            pl.lit(snapshot_date).alias("date"),
            pl.col("expiration").str.strptime(pl.Date, "%d/%m/%y"),
            pl.col("strike").cast(pl.Float64),
            pl.col("contract_size").cast(pl.Int64),
        )
        .pipe(_add_symbol, symbol)
        .select(sorted(_CHAIN_COLS))
        .sort(["expiration", "strike"])
    )