# Sync flush: drain bars from Rust ring buffer → ClickHouse write (Issue #318)
from __future__ import annotations

import logging
from collections import defaultdict

import polars as pl

from .arrow_flush import _sync_flush_rest_bars_arrow
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__)


def _sync_flush_rest_bars(
    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:
    """Dispatch to Arrow or dict drain path. See _sync_flush_rest_bars_arrow/dict."""
    if hasattr(engine, "drain_bars_arrow"):
        return _sync_flush_rest_bars_arrow(
            engine, valid_pairs, symbol_modes=symbol_modes, cache=cache,
            rust_ch_writer_mode=rust_ch_writer_mode,
        )
    return _sync_flush_rest_bars_dict(
        engine, valid_pairs, symbol_modes=symbol_modes, cache=cache,
        rust_ch_writer_mode=rust_ch_writer_mode,
    )


def _sync_flush_rest_bars_dict(  # noqa: PLR0915  INLINE-IGNORE-OK
    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:
    """Dict-based sync_flush: legacy fallback for OdbEngine or older builds (#296)."""
    if hasattr(engine, "drain_bars"):
        all_bars: list[dict] = engine.drain_bars()
    else:
        # Fallback for OdbEngine or older builds without drain_bars
        all_bars = []
        while True:
            bar = engine.next_bar(timeout_ms=500)
            if bar is None:
                break
            all_bars.append(bar)

    raw_count = len(all_bars)
    batches: dict[tuple[str, int], list[dict]] = defaultdict(list)
    drained = 0
    _seen_rejected: set[str] = set()

    if valid_pairs is not None:
        vp_sample = sorted(valid_pairs)[:5]
        vp_types = {
            (type(s).__name__, type(t).__name__) for s, t in list(valid_pairs)[:3]
        }
        logger.info(
            "sync_flush: valid_pairs=%d, sample=%s, types=%s",
            len(valid_pairs),
            vp_sample,
            vp_types,
        )

    for bar in all_bars:
        symbol = bar.pop("_symbol", "UNKNOWN")
        threshold = bar.pop("_threshold", 0)

        # Skip pairs not in valid_pairs
        if valid_pairs is not None and (symbol, threshold) not in valid_pairs:
            _pair_key = f"{symbol}@{threshold}"
            if _pair_key not in _seen_rejected:
                _seen_rejected.add(_pair_key)
                logger.info(
                    "sync_flush: REJECTED %s (sym_type=%s, thr_type=%s, thr_repr=%r)",
                    _pair_key,
                    type(symbol).__name__,
                    type(threshold).__name__,
                    threshold,
                )
            continue

        batches[(symbol, threshold)].append(bar)
        drained += 1

    # MEM: Release the raw drain list — bars are now in per-pair batches.
    # This frees ~800MB of dict references for a 200K-bar drain.
    del all_bars

    filtered_invalid_pairs = raw_count - drained
    logger.info(
        "sync_flush: drain diagnostics: raw=%d, valid=%d, invalid_pair=%d",
        raw_count,
        drained,
        filtered_invalid_pairs,
    )

    if drained == 0:
        logger.info(
            "sync_flush: ring buffer empty — nothing to flush (raw drain: %d bars)",
            raw_count,
        )
        return 0

    # --- Stathera continuity check within each (symbol, threshold) batch ---
    # Diagnostic only: gaps and non-monotonic ordering are logged but never
    # block the CH write. fill_from_rest should produce gapless bars, but
    # ring buffer overflow or edge cases could introduce discontinuities.
    for (sym, thr), batch_bars in batches.items():
        if len(batch_bars) < 2:  # noqa: PLR2004
            continue
        # LAZY-02: In-place sort avoids copying the entire batch list.
        batch_bars.sort(key=lambda b: b.get("first_agg_trade_id", 0))
        for i in range(1, len(batch_bars)):
            prev_last = batch_bars[i - 1].get("last_agg_trade_id", 0)
            curr_first = batch_bars[i].get("first_agg_trade_id", 0)
            if curr_first <= prev_last:
                logger.warning(
                    "sync_flush: stathera non-monotonic in REST batch %s@%d: "
                    "bar[%d].last_tid=%d >= bar[%d].first_tid=%d",
                    sym,
                    thr,
                    i - 1,
                    prev_last,
                    i,
                    curr_first,
                )
            elif curr_first != prev_last + 1:
                logger.warning(
                    "sync_flush: stathera gap in REST batch %s@%d: "
                    "bar[%d].last_tid=%d -> bar[%d].first_tid=%d (delta=%d)",
                    sym,
                    thr,
                    i - 1,
                    prev_last,
                    i,
                    curr_first,
                    curr_first - prev_last - 1,
                )

    from opendeviationbar import __version__
    from opendeviationbar.clickhouse.cache import OpenDeviationBarCache

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

    # ALLOC-01/03: Use shared cache and pre-computed symbol_modes
    _modes = symbol_modes or {}
    _cache = cache
    if _cache is None:
        _cache = OpenDeviationBarCache()

    for (symbol, threshold), bars in batches.items():
        _mode = _modes.get(symbol, "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 += len(bars)
            continue
        try:
            bars_df = pl.DataFrame(bars)
            inserted = _cache.store_bars_batch(
                symbol=symbol,
                threshold_decimal_bps=threshold,
                bars=bars_df,
                version=source_version,
                overlap_strategy="off",
                ouroboros_mode=_mode,
                _skip_lock=False,
            )
            written += inserted
        except Exception:
            logger.info(
                "sync_flush: store_bars_batch failed for %s@%d, "
                "falling back to _write_batch_with_retry (non-fatal)",
                symbol,
                threshold,
                exc_info=True,
            )
            if _write_batch_with_retry(symbol, threshold, bars, cache=_cache, ouroboros_mode=_mode):
                written += len(bars)

    # Populate junction telemetry for /status endpoint
    for (symbol, _threshold), bars in batches.items():
        if symbol not in _junction_state:
            _junction_state[symbol] = {}
        last_tid = max(
            (b.get("last_agg_trade_id", 0) for b in bars),
            default=0,
        )
        if last_tid > 0:
            _junction_state[symbol]["rest_last_tid"] = last_tid
        min_tid = min(
            (b.get("first_agg_trade_id", 0) for b in bars),
            default=0,
        )
        if min_tid > 0:
            _junction_state[symbol]["rest_first_tid"] = min(
                _junction_state[symbol].get("rest_first_tid", min_tid),
                min_tid,
            )

    # Diagnostic: log which pairs were written and which had zero bars
    pairs_with_bars = sorted(
        ((s, t, len(bars)) for (s, t), bars in batches.items()),
        key=lambda x: (-x[2], x[0], x[1]),
    )
    logger.info(
        "sync_flush: %d/%d REST bars written to CH before start() (%d pairs)",
        written,
        drained,
        len(batches),
    )
    _max_pairs_to_log = 50
    if len(pairs_with_bars) <= _max_pairs_to_log:
        for sym, thr, cnt in pairs_with_bars:
            logger.debug("sync_flush: wrote %s@%d = %d bars", sym, thr, cnt)

    return written
