"""Quality assurance on raw bars (never alters prices; produces audit tables and row masks)."""
from __future__ import annotations

import numpy as np
import polars as pl


def qa_report(bars: pl.DataFrame) -> dict:
    """Row-level integrity checks on a bars frame (one or more symbols). Returns counts, never mutates."""
    n = bars.height
    if n == 0:
        return {"rows": 0}
    ohlc_bad = bars.filter((pl.col("high") < pl.col("low")) | (pl.col("open") > pl.col("high")) |
                           (pl.col("open") < pl.col("low")) | (pl.col("close") > pl.col("high")) |
                           (pl.col("close") < pl.col("low"))).height
    nonpos = bars.filter((pl.col("open") <= 0) | (pl.col("high") <= 0) | (pl.col("low") <= 0) |
                         (pl.col("close") <= 0)).height
    negvol = bars.filter(pl.col("volume") < 0).height
    zerovol = bars.filter(pl.col("volume") == 0).height
    dup = n - bars.select(["instrument_id", "timestamp_utc"]).unique().height
    offint = bars.filter(pl.col("timestamp_utc").dt.second() != 0).height
    flagged = bars.filter(pl.col("quality_flags").list.len() > 0).height
    revs = bars.select(pl.col("revision_id").unique()).to_series().to_list()
    return {"rows": n, "ohlc_inconsistent": ohlc_bad, "nonpositive_price": nonpos, "negative_volume": negvol,
            "zero_volume": zerovol, "duplicates": dup, "off_interval": offint, "flagged_rows": flagged,
            "revision_ids": sorted(int(r) for r in revs), "symbols": bars["symbol_as_reported"].n_unique(),
            "first_ts": str(bars["timestamp_utc"].min()), "last_ts": str(bars["timestamp_utc"].max())}


def dedupe(bars: pl.DataFrame) -> pl.DataFrame:
    """Keep the highest revision per (instrument, timestamp). The lake may publish the last session twice (base+delta)."""
    return (bars.sort(["instrument_id", "timestamp_utc", "revision_id"])
            .unique(subset=["instrument_id", "timestamp_utc"], keep="last", maintain_order=True))


def price_jump_mask(close: np.ndarray, valid: np.ndarray, thresh: float = 0.25) -> np.ndarray:
    """Flag consecutive valid closes whose absolute log change exceeds `thresh` (vendor code: price_jump)."""
    out = np.zeros_like(valid, dtype=bool)
    idx = np.flatnonzero(valid)
    if len(idx) < 2:
        return out
    r = np.abs(np.diff(np.log(close[idx])))
    out[idx[1:][r > thresh]] = True
    return out


def session_coverage(bars: pl.DataFrame, expected: pl.DataFrame) -> pl.DataFrame:
    """Per (symbol, session): regular-hours bars present / expected slots. `expected` has session_date, n_slots."""
    reg = bars.filter(pl.col("slot") >= 0).group_by(["symbol_as_reported", "session_date"]).agg(
        pl.len().alias("present"), pl.col("volume").sum().alias("volume"))
    cov = reg.join(expected, on="session_date", how="left")
    return cov.with_columns((pl.col("present") / pl.col("n_slots")).alias("coverage")).sort(
        ["symbol_as_reported", "session_date"])