# Issue #313: Arrow-based sync_flush for zero-copy drain + Polars filter (v7.2)
# Issue #318: Rust CH writer feature flag gating
"""Arrow FFI drain path for _sync_flush_rest_bars.

Drains ring buffer as Arrow RecordBatch with _symbol and _threshold_decimal_bps
metadata columns. Uses Polars df.filter() for zero-copy per-pair slicing instead
of partition_by (which copies data). Eliminates ~200K Python dict allocations.
"""
from __future__ import annotations

import logging

import polars as pl

from .ch_writer_flag import _should_python_write
from .clickhouse_write import _write_batch_with_retry
from .http_server import _junction_state

logger = logging.getLogger(__name__)

# Track verified TID per (symbol, threshold) — skip re-checking already-verified ranges
_verified_up_to: dict[tuple[str, int], int] = {}


def _verify_boundary(
    cache: object,
    symbol: str,
    threshold: int,
    pair_df: pl.DataFrame,
) -> None:
    """Incremental boundary Stathera check after a batch write (#318).

    Queries the single bar immediately BEFORE this batch's min(first_agg_trade_id)
    and verifies ``prev.last_agg_trade_id + 1 == batch.min_first``. Single-row
    query, ~2-3ms. Skips if the range was already verified (``_verified_up_to``).
    """
    try:
        min_first = int(pair_df["first_agg_trade_id"].min())
        max_last = int(pair_df["last_agg_trade_id"].max())

        # Skip if this range was already verified
        key = (symbol, threshold)
        prev_verified = _verified_up_to.get(key, 0)
        if min_first <= prev_verified:
            return  # Already verified up to this point

        result = cache.client.query(
            "SELECT last_agg_trade_id"
            " FROM opendeviationbar_cache.open_deviation_bars FINAL"
            f" WHERE symbol = '{symbol}'"
            f" AND threshold_decimal_bps = {threshold}"
            f" AND last_agg_trade_id < {min_first}"
            " AND first_agg_trade_id > 0"
            " ORDER BY last_agg_trade_id DESC"
            " LIMIT 1"
        )

        if not result.result_rows:
            # No prior bar — this is the first batch for this pair (clean)
            _verified_up_to[key] = max_last
            return

        prev_last_tid = int(result.result_rows[0][0])
        expected_first = prev_last_tid + 1
        delta = min_first - expected_first

        if delta == 0:
            # Clean boundary
            _verified_up_to[key] = max_last
        elif delta > 0:
            logger.warning(
                "stathera-boundary[%s@%d]: GAP detected — "
                "prev_last_tid=%d, batch_min_first=%d, delta=%d trades missing",
                symbol, threshold, prev_last_tid, min_first, delta,
            )
        else:
            logger.critical(
                "stathera-boundary[%s@%d]: OVERLAP detected — "
                "prev_last_tid=%d, batch_min_first=%d, delta=%d trades overlapping",
                symbol, threshold, prev_last_tid, min_first, abs(delta),
            )
    except Exception:
        # Fail-open: verification failure must never block writes
        logger.debug("stathera-boundary check failed (non-fatal)", exc_info=True)


def _sync_flush_rest_bars_arrow(
    engine: object,
    valid_pairs: set[tuple[str, int]] | None,
    *,
    symbol_modes: dict[str, str] | None = None,
    cache: object | None = None,
    rust_ch_writer_mode: str = "off",
) -> int:
    """Arrow-based sync_flush: zero-copy drain + Polars filter (#313, v7.2).

    Called by _sync_flush_rest_bars when drain_bars_arrow() is available.
    Uses df.filter() for zero-copy per-pair slicing (no data duplication).

    Parameters
    ----------
    engine
        StreamManager with drain_bars_arrow() method.
    valid_pairs
        Set of (symbol, threshold) pairs to accept. None = accept all.
    symbol_modes
        Per-symbol ouroboros mode ("aion" or "day").

    Returns
    -------
    int
        Number of bars written to ClickHouse.
    """

    # Drain as Arrow RecordBatch -- GIL released during Rust work
    arrow_batch = engine.drain_bars_arrow()

    # Zero-copy convert to Polars DataFrame
    df: pl.DataFrame = pl.from_arrow(arrow_batch)
    raw_count = len(df)

    if raw_count == 0:
        logger.info("sync_flush[arrow]: ring buffer empty -- nothing to flush")
        return 0

    logger.info("sync_flush[arrow]: drained %d bars as Arrow RecordBatch", raw_count)

    # --- Filter: valid_pairs (semi-join) ---
    if valid_pairs is not None:
        vp_sample = sorted(valid_pairs)[:5]
        logger.info(
            "sync_flush[arrow]: valid_pairs=%d, sample=%s",
            len(valid_pairs),
            vp_sample,
        )
        # Build a filter DataFrame from valid_pairs for semi-join
        vp_df = pl.DataFrame(
            {
                "_symbol": [s for s, _ in valid_pairs],
                "_threshold_decimal_bps": [
                    t if isinstance(t, int) else int(t) for _, t in valid_pairs
                ],
            },
            schema={
                "_symbol": pl.Utf8,
                "_threshold_decimal_bps": pl.UInt32,
            },
        )
        df = df.join(vp_df, on=["_symbol", "_threshold_decimal_bps"], how="semi")  # noqa: PD901

    drained = len(df)
    filtered_invalid_pairs = raw_count - drained
    logger.info(
        "sync_flush[arrow]: diagnostics: raw=%d, valid=%d, invalid_pair=%d",
        raw_count,
        drained,
        filtered_invalid_pairs,
    )

    if drained == 0:
        logger.info("sync_flush[arrow]: no valid bars after filtering")
        return 0

    # --- Stathera continuity diagnostic per (symbol, threshold) ---
    _stathera_diagnostic_polars(df)

    # --- Write to CH per (symbol, threshold) using df.filter() (zero-copy) ---
    from opendeviationbar import __version__
    from opendeviationbar.clickhouse.cache import OpenDeviationBarCache

    source_version = f"{__version__}+sidecar"
    written = 0

    # ALLOC-03: Use pre-computed symbol_modes instead of per-flush registry lookup
    _modes = symbol_modes or {}

    # Get unique (symbol, threshold) pairs via Polars -- no partition_by
    pairs = (
        df.select("_symbol", "_threshold_decimal_bps")
        .unique()
        .iter_rows()
    )

    # ALLOC-01: Use shared cache if provided, else create one
    _cache = cache if cache is not None else OpenDeviationBarCache()

    for sym, thr in pairs:
        # Zero-copy filter -- returns a view on Arrow-backed data
        pair_df = df.filter(
            (pl.col("_symbol") == sym)
            & (pl.col("_threshold_decimal_bps") == thr)
        ).drop("_symbol", "_threshold_decimal_bps")

        threshold = int(thr)
        _mode = _modes.get(sym, "aion")

        # Issue #318: Gate Python CH write based on feature flag
        if not _should_python_write(rust_ch_writer_mode):
            # live mode: Rust already wrote -- just count bars
            written += pair_df.height
            continue

        try:
            inserted = _cache.store_bars_batch(
                symbol=sym,
                threshold_decimal_bps=threshold,
                bars=pair_df,
                version=source_version,
                overlap_strategy="off",
                ouroboros_mode=_mode,
                _skip_lock=False,
            )
            written += inserted

            # Incremental boundary verification (#318): after successful write,
            # check that the bar immediately before this batch connects cleanly.
            # Single-row query (~2-3ms), catches gaps/overlaps at write time.
            if inserted > 0 and _cache is not None:
                _verify_boundary(
                    _cache, sym, threshold, pair_df,
                )
        except Exception:
            logger.info(
                "sync_flush[arrow]: store_bars_batch failed for %s@%d, "
                "falling back to _write_batch_with_retry (non-fatal)",
                sym,
                threshold,
                exc_info=True,
            )
            # Fallback: convert pair_df to list[dict] for retry path
            bar_dicts = pair_df.to_dicts()
            if _write_batch_with_retry(
                sym, threshold, bar_dicts, cache=_cache, ouroboros_mode=_mode
            ):
                written += len(bar_dicts)

    # Populate junction telemetry for /status endpoint (Polars aggregation)
    for sym in df["_symbol"].unique().to_list():
        sym_df = df.filter(pl.col("_symbol") == sym)
        if sym not in _junction_state:
            _junction_state[sym] = {}
        last_tid = sym_df["last_agg_trade_id"].max()
        if last_tid and last_tid > 0:
            _junction_state[sym]["rest_last_tid"] = last_tid
        min_tid = sym_df["first_agg_trade_id"].min()
        if min_tid and min_tid > 0:
            _junction_state[sym]["rest_first_tid"] = min(
                _junction_state[sym].get("rest_first_tid", min_tid),
                min_tid,
            )

    # Diagnostic: log pair counts
    pairs_with_bars = (
        df.group_by("_symbol", "_threshold_decimal_bps")
        .len()
        .sort("len", descending=True)
    )
    logger.info(
        "sync_flush[arrow]: %d/%d REST bars written to CH (Arrow path, %d pairs)",
        written,
        drained,
        pairs_with_bars.height,
    )

    return written


def _stathera_diagnostic_polars(df: pl.DataFrame) -> None:
    """Stathera continuity check using Polars operations (#313, v7.2).

    Uses df.filter() for zero-copy per-pair slicing and Polars sort
    instead of Python sorted(). Diagnostic only: gaps and non-monotonic
    ordering are logged but never block the CH write.
    """
    pairs = (
        df.select("_symbol", "_threshold_decimal_bps")
        .unique()
        .iter_rows()
    )

    for sym, thr in pairs:
        pair_df = df.filter(
            (pl.col("_symbol") == sym)
            & (pl.col("_threshold_decimal_bps") == thr)
        )
        if pair_df.height < 2:  # noqa: PLR2004
            continue

        sorted_df = pair_df.sort("first_agg_trade_id")
        prev_last = sorted_df["last_agg_trade_id"].shift(1)
        curr_first = sorted_df["first_agg_trade_id"]

        # Skip first row (no previous bar) via slice
        deltas = (curr_first - prev_last).slice(1)

        non_monotonic = (deltas <= 0).sum()
        if non_monotonic > 0:
            logger.warning(
                "sync_flush[arrow]: stathera non-monotonic in REST batch "
                "%s@%d: %d occurrences",
                sym,
                int(thr),
                non_monotonic,
            )

        gaps = ((deltas > 1) | (deltas < 0)).sum()
        if gaps > 0:
            logger.warning(
                "sync_flush[arrow]: stathera gaps in REST batch "
                "%s@%d: %d occurrences",
                sym,
                int(thr),
                gaps,
            )
