"""Row-level and session-level validation of canonical bars. Flags, never mutates prices."""
from __future__ import annotations

from datetime import date
from typing import Any

import polars as pl
import pyarrow as pa
from findata_schema import INTERVAL_SECONDS, Interval, QualityFlag

from findata_quality.calendar import SessionCalendar, xnys

PRICE_JUMP_THRESHOLD = 0.25  # 25 % between consecutive 1-minute closes inside a session (flag only)


NO_SESSION_CALENDAR = {"fx", "crypto", "futures"}  # 24h / overnight sessions: no XNYS session checks


def validate_bars(table: pa.Table, interval: Interval, calendar: SessionCalendar | None = None,
                  check_sessions: bool = True, asset_class: str | None = None) -> tuple[pa.Table, list[dict[str, Any]]]:
    """Return (table with `quality_flags` filled, anomalies list). Input must follow BARS_SCHEMA columns."""
    cal = calendar or xnys()
    step = INTERVAL_SECONDS[interval]
    df = pl.from_arrow(table)
    if df.height == 0:
        return table, []
    if asset_class is None and "asset_class" in df.columns:
        asset_class = df.get_column("asset_class")[0]
    if asset_class in NO_SESSION_CALENDAR:
        check_sessions = False
    df = df.sort(["symbol_as_reported", "timestamp_utc"])
    ts_s = pl.col("timestamp_utc").dt.epoch("s")
    flags: list[pl.Expr] = [
        pl.when((pl.col("high") < pl.max_horizontal("open", "close")) | (pl.col("low") > pl.min_horizontal("open", "close"))
                | (pl.col("high") < pl.col("low"))).then(pl.lit(QualityFlag.ohlc_inconsistent.value)),
        pl.when(pl.min_horizontal("open", "high", "low", "close") <= 0).then(pl.lit(QualityFlag.nonpositive_price.value)),
        pl.when(pl.col("volume") < 0).then(pl.lit(QualityFlag.negative_volume.value)),
        pl.when((pl.col("volume") == 0) & pl.lit(asset_class != "index")).then(pl.lit(QualityFlag.zero_volume.value)),
        pl.when(pl.col("volume") != pl.col("volume").floor()).then(pl.lit(QualityFlag.fractional_volume.value)),
        pl.when((ts_s % step) != 0).then(pl.lit(QualityFlag.off_interval.value)),
        pl.when(pl.struct(["symbol_as_reported", "timestamp_utc"]).is_duplicated()).then(pl.lit(QualityFlag.duplicate.value)),
    ]
    if interval != Interval.d1:
        prev_close = pl.col("close").shift(1).over(["symbol_as_reported", "session_date"])
        flags.append(pl.when(((pl.col("close") / prev_close) - 1).abs() > PRICE_JUMP_THRESHOLD)
                     .then(pl.lit(QualityFlag.price_jump.value)))
    if check_sessions and interval != Interval.d1:
        # outside extended hours / non trading day: computed per distinct session_date
        dates = df.get_column("session_date").unique().to_list()
        bounds = {d: cal.extended_bounds_utc(d) for d in dates}
        non_sessions = [d for d in dates if not cal.is_session(d)]
        lo = pl.col("session_date").replace_strict({d: b[0] for d, b in bounds.items()}, return_dtype=pl.Datetime("us", "UTC"))
        hi = pl.col("session_date").replace_strict({d: b[1] for d, b in bounds.items()}, return_dtype=pl.Datetime("us", "UTC"))
        flags.append(pl.when((pl.col("timestamp_utc") < lo) | (pl.col("timestamp_utc") >= hi))
                     .then(pl.lit(QualityFlag.outside_session.value)))
        if non_sessions:
            flags.append(pl.when(pl.col("session_date").is_in(non_sessions)).then(pl.lit(QualityFlag.non_trading_day.value)))
    df = df.with_columns(
        pl.concat_list([f.otherwise(None) for f in flags]).list.drop_nulls().alias("quality_flags")
    )
    anomalies = _summarize(df)
    out = df.to_arrow()
    # cast back to the canonical schema (polars may widen list/int types)
    out = out.select(table.column_names).cast(table.schema)
    return out, anomalies


def _summarize(df: pl.DataFrame) -> list[dict[str, Any]]:
    flagged = df.filter(pl.col("quality_flags").list.len() > 0)
    if flagged.height == 0:
        return []
    agg = (flagged.explode("quality_flags")
           .group_by(["symbol_as_reported", "session_date", "quality_flags"]).agg(pl.len().alias("n"))
           .sort(["symbol_as_reported", "session_date", "quality_flags"]))
    sev = {QualityFlag.zero_volume.value: "info", QualityFlag.fractional_volume.value: "info",
           QualityFlag.price_jump.value: "warning", QualityFlag.outside_session.value: "warning",
           QualityFlag.gap.value: "info"}
    return [{"symbol": r["symbol_as_reported"], "session_date": r["session_date"], "code": r["quality_flags"],
             "severity": sev.get(r["quality_flags"], "error"), "detail": {"rows": r["n"]}}
            for r in agg.iter_rows(named=True)]


def session_coverage(table: pa.Table, interval: Interval, calendar: SessionCalendar | None = None) -> list[dict[str, Any]]:
    """Per (symbol, session_date): bar count, first/last ts, expected regular-hours bars, gap flags."""
    cal = calendar or xnys()
    step = INTERVAL_SECONDS[interval]
    df = pl.from_arrow(table)
    if df.height == 0:
        return []
    g = (df.group_by(["symbol_as_reported", "session_date"])
         .agg(pl.len().alias("bar_count"), pl.col("timestamp_utc").min().alias("first_ts"),
              pl.col("timestamp_utc").max().alias("last_ts"))
         .sort(["symbol_as_reported", "session_date"]))
    rows = []
    for r in g.iter_rows(named=True):
        d: date = r["session_date"]
        exp = cal.expected_regular_bars(d, step) if interval != Interval.d1 else 1
        flags = []
        if not cal.is_session(d):
            flags.append(QualityFlag.non_trading_day.value)
        rows.append({"symbol": r["symbol_as_reported"], "session_date": d, "bar_count": r["bar_count"],
                     "first_ts": r["first_ts"], "last_ts": r["last_ts"], "expected_bars": exp, "flags": flags})
    return rows