from __future__ import annotations

from datetime import datetime
from typing import Any

import polars as pl

from uforia.storage import DatasetName, DatasetStore

UTC_DT = pl.Datetime(time_zone="UTC")
_HEDGE_ROUTE_SCHEMA: dict[str, Any] = {
    "ts": UTC_DT,
    "spot_mid": pl.Float64,
    "futures_mid": pl.Float64,
    "perp_mid": pl.Float64,
    "futures_basis_bps": pl.Float64,
    "perp_basis_bps": pl.Float64,
    "futures_perp_spread_bps": pl.Float64,
    "route_basis_bps": pl.Float64,
    "micro_route_score": pl.Float64,
    "route_preference": pl.String,
}


def session_bucket_expr(ts_column: str = "ts") -> pl.Expr:
    hour = pl.col(ts_column).dt.hour()
    return (
        pl.when(hour < 8)
        .then(pl.lit("asia"))
        .when(hour < 16)
        .then(pl.lit("europe"))
        .otherwise(pl.lit("us"))
    )


def rolling_corr_expr(left: str, right: str, window: int, min_samples: int | None = None) -> pl.Expr:
    samples = min_samples or max(12, window // 8)
    x = pl.col(left)
    y = pl.col(right)
    cov = (x * y).rolling_mean(window_size=window, min_samples=samples) - (
        x.rolling_mean(window_size=window, min_samples=samples)
        * y.rolling_mean(window_size=window, min_samples=samples)
    )
    left_std = x.rolling_std(window_size=window, min_samples=samples)
    right_std = y.rolling_std(window_size=window, min_samples=samples)
    return (cov / ((left_std * right_std) + 1e-9)).clip(-1.0, 1.0)


def build_hedge_route_frame(
    store: DatasetStore,
    *,
    start: datetime,
    end: datetime,
    asset: str,
) -> pl.DataFrame:
    bars = store.read(
        DatasetName.MICROSTRUCTURE_BAR,
        start=start,
        end=end,
        filters={"asset": asset},
    )
    if bars.is_empty():
        return pl.DataFrame(schema=_HEDGE_ROUTE_SCHEMA)

    preferred_bar_size = bars.select(pl.col("bar_size_secs").drop_nulls().max()).item()
    if preferred_bar_size is not None:
        bars = bars.filter(pl.col("bar_size_secs") == preferred_bar_size)
    if bars.is_empty():
        return pl.DataFrame(schema=_HEDGE_ROUTE_SCHEMA)

    market_type = pl.coalesce([pl.col("market_type"), _market_type_expr("venue")])
    contract_type = pl.coalesce([pl.col("contract_type"), _contract_type_expr("venue")])
    normalized = bars.with_columns(
        market_type.alias("market_type"),
        contract_type.alias("contract_type"),
    )
    aggregated = (
        normalized.sort(["ts", "venue"])
        .group_by("ts")
        .agg(
            pl.col("close").filter(pl.col("market_type") == "spot").mean().alias("spot_mid"),
            pl.col("close")
            .filter(
                (pl.col("market_type") == "futures")
                | (pl.col("contract_type") == "linear_future")
            )
            .mean()
            .alias("futures_mid"),
            pl.col("close")
            .filter((pl.col("market_type") == "perp") | (pl.col("contract_type") == "linear_perp"))
            .mean()
            .alias("perp_mid"),
            pl.col("spread")
            .filter(
                (pl.col("market_type") == "futures")
                | (pl.col("contract_type") == "linear_future")
            )
            .mean()
            .alias("futures_spread_abs"),
            pl.col("spread")
            .filter((pl.col("market_type") == "perp") | (pl.col("contract_type") == "linear_perp"))
            .mean()
            .alias("perp_spread_abs"),
        )
        .sort("ts")
    )
    if aggregated.is_empty():
        return pl.DataFrame(schema=_HEDGE_ROUTE_SCHEMA)

    denominator = pl.coalesce(
        [
            pl.col("spot_mid"),
            pl.col("perp_mid"),
            pl.col("futures_mid"),
            pl.lit(1e-9),
        ]
    )
    routed = aggregated.with_columns(
        (((pl.col("futures_mid") - pl.col("spot_mid")) / denominator) * 10_000.0).alias(
            "futures_basis_bps"
        ),
        (((pl.col("perp_mid") - pl.col("spot_mid")) / denominator) * 10_000.0).alias(
            "perp_basis_bps"
        ),
        (
            ((pl.col("futures_mid") - pl.col("perp_mid")).abs() / denominator) * 10_000.0
        ).alias("futures_perp_spread_bps"),
    ).with_columns(
        pl.when(pl.col("futures_mid").is_not_null())
        .then(
            100.0
            - pl.col("futures_basis_bps").fill_null(0.0).abs()
            - (pl.col("futures_spread_abs").fill_null(0.0) * 10.0)
            - (pl.col("futures_perp_spread_bps").fill_null(0.0) * 0.25)
        )
        .otherwise(None)
        .alias("future_route_score"),
        pl.when(pl.col("perp_mid").is_not_null())
        .then(
            100.0
            - pl.col("perp_basis_bps").fill_null(0.0).abs()
            - (pl.col("perp_spread_abs").fill_null(0.0) * 10.0)
            - (pl.col("futures_perp_spread_bps").fill_null(0.0) * 0.10)
        )
        .otherwise(None)
        .alias("perp_route_score"),
    ).with_columns(
        pl.when(
            pl.col("future_route_score").is_not_null()
            & (
                pl.col("perp_route_score").is_null()
                | (pl.col("future_route_score") >= pl.col("perp_route_score"))
            )
        )
        .then(pl.lit("nearby_future"))
        .when(pl.col("perp_route_score").is_not_null())
        .then(pl.lit("perp"))
        .otherwise(None)
        .alias("route_preference")
    ).with_columns(
        pl.when(pl.col("route_preference") == "nearby_future")
        .then(pl.col("future_route_score"))
        .when(pl.col("route_preference") == "perp")
        .then(pl.col("perp_route_score"))
        .otherwise(None)
        .alias("micro_route_score"),
        pl.when(pl.col("route_preference") == "nearby_future")
        .then(pl.col("futures_basis_bps").abs())
        .when(pl.col("route_preference") == "perp")
        .then(pl.col("perp_basis_bps").abs())
        .otherwise(None)
        .alias("route_basis_bps"),
    )
    return routed.select(_HEDGE_ROUTE_SCHEMA.keys())


def _market_type_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("spot"))
        .then(pl.lit("spot"))
        .when(pl.col(column).str.contains("futures"))
        .then(pl.lit("futures"))
        .when(pl.col(column).str.contains("perps"))
        .then(pl.lit("perp"))
        .otherwise(pl.lit("perp"))
    )


def _contract_type_expr(column: str) -> pl.Expr:
    return (
        pl.when(pl.col(column).str.contains("spot"))
        .then(pl.lit("spot"))
        .when(pl.col(column).str.contains("futures"))
        .then(pl.lit("linear_future"))
        .otherwise(pl.lit("linear_perp"))
    )
