#!/usr/bin/env python3
"""Direct parallel Parquet->Rust->ClickHouse repair (Issue #275).

Bypasses kintsugi's shard discovery overhead. Reads raw trades from the
local Parquet tick cache, processes through the Rust processor (with the
v13.20.0 fix for orphan microstructure), and writes bars to ClickHouse.

Processes ALL 4 thresholds per day from a single Parquet read -- the tick
cache is threshold-agnostic (raw aggTrades), so one read fans out to 4
processors. This is 4x more efficient than kintsugi which processes
each (symbol, threshold) shard independently.

FD-safety (Issue #279): Uses ONE shared OpenDeviationBarCache per worker
thread (via threading.local), not per-threshold-per-day. This reduces
concurrent CH connections from 4*N to N, preventing EMFILE with the
default ulimit -n 1024 on bigblack.

Memory-safety:
  - Dict path: Uses .to_dicts() for Polars->Rust interop. Arrow zero-copy
    (process_trades_arrow) was planned but deferred due to column name/type
    parity issues between Arrow output and store_bars_batch() expectations.
  - MEM-018: Bounded future submission (max 2*workers pending futures)
    instead of submitting all 2,000+ days at once. Completed futures are
    discarded immediately instead of being held in a dict until loop exit.
  - Allocator purge: _force_allocator_purge() every 50 completions returns
    freed glibc/mimalloc pages to the OS, preventing RSS plateau from
    arena fragmentation.
  - overlap_strategy="warn": Intentional -- plain INSERT avoids the 300s
    DELETE scan timeout that overlap_strategy="replace" causes when
    scanning 1,658+ active parts. Safe because repair targets days with
    no existing bars.
  Expected RSS: ~3 GB start, ~13 GB steady state with periodic purge.

Usage:
    # Dry run
    python scripts/repair_direct_parquet.py

    # Execute (BTCUSDT, 24 workers)
    python scripts/repair_direct_parquet.py --execute --workers 24

    # Single symbol
    python scripts/repair_direct_parquet.py --execute --symbol BTCUSDT
"""
from __future__ import annotations

import argparse
import gc
import logging
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, datetime, timedelta, timezone
from itertools import islice as _islice

# Module-level imports for hot path (avoid per-call dict lookup overhead)
import polars as pl
from opendeviationbar import __version__ as _odb_version
from opendeviationbar._core import PyOpenDeviationBarProcessor as _Processor
from opendeviationbar.storage.parquet import TickStorage as _TickStorage

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# Bounded submission: keep at most 3*workers futures pending.
# Higher multiplier keeps the pipeline fuller — workers never starve
# waiting for the main thread to submit the next batch.
_PENDING_MULTIPLIER = 3

# Purge allocator arenas every N completed days to prevent RSS bloat.
# With the Arrow zero-copy path, memory pressure is much lower (~30 MB/day
# vs ~150 MB/day with dicts), so purge less frequently.
_PURGE_INTERVAL = 100

THRESHOLDS_DEFAULT = [250, 500, 750]

# Column selection constant (avoid per-call list allocation)
_TICK_COLS = ["agg_trade_id", "price", "quantity", "timestamp", "is_buyer_maker"]
_TICK_COLS_SET = frozenset(_TICK_COLS)

# µs detection threshold: legitimate ms timestamp cannot reach this until year 47658
_US_THRESHOLD = 1_500_000_000_000_000

# Pre-computed source version string (avoid f-string allocation per day per threshold)
_SOURCE_VERSION = f"{_odb_version}+repair-parquet"

# Thread-local storage for per-worker CH cache reuse.
# Each worker thread gets ONE OpenDeviationBarCache for its lifetime,
# instead of creating 4 per day (which was the root cause of fd exhaustion).
_thread_local = threading.local()

# Shared TickStorage instance (stateless, thread-safe)
_shared_storage = None


def _get_worker_cache() -> object:
    """Get or create a thread-local OpenDeviationBarCache.

    Returns the same cache instance for repeated calls within the same
    thread. This eliminates the per-threshold connection overhead:
    - No repeated preflight checks
    - No repeated _ensure_schema() migrations
    - No new urllib3.PoolManager per call (each holds up to 20 fds)
    """
    cache = getattr(_thread_local, "cache", None)
    if cache is None:
        from opendeviationbar.clickhouse.cache import OpenDeviationBarCache

        cache = OpenDeviationBarCache()
        _thread_local.cache = cache
        logger.debug("Created thread-local CH cache for %s", threading.current_thread().name)
    return cache


def _close_worker_cache() -> None:
    """Close and discard the thread-local cache (called at shutdown)."""
    import contextlib

    cache = getattr(_thread_local, "cache", None)
    if cache is not None:
        with contextlib.suppress(Exception):
            cache.close()
        _thread_local.cache = None


_purge_lib = None  # Cached ctypes handle (avoid dlopen per call)
_purge_method = None  # "mimalloc" | "glibc" | "none"


def _force_allocator_purge() -> None:
    """Force mimalloc/glibc to return freed pages to the OS immediately.

    Caches the ctypes library handle on first call — subsequent calls
    are a single FFI function invocation with zero allocation overhead.
    """
    import sys

    global _purge_lib, _purge_method  # noqa: PLW0603

    if sys.platform != "linux":
        return

    # First call: resolve which allocator to use
    if _purge_method is None:
        import ctypes

        try:
            _purge_lib = ctypes.CDLL("libmimalloc.so.2")
            _purge_method = "mimalloc"
        except OSError:
            try:
                _purge_lib = ctypes.CDLL("libc.so.6")
                _purge_method = "glibc"
            except OSError:
                _purge_method = "none"

    if _purge_method == "mimalloc":
        import ctypes

        _purge_lib.mi_collect(ctypes.c_bool(True))
    elif _purge_method == "glibc":
        _purge_lib.malloc_trim(0)


def _get_days_needing_repair(symbol: str, thresholds: list[int]) -> list[date]:
    """Find days missing from ClickHouse that have Parquet ticks."""
    import clickhouse_connect

    client = clickhouse_connect.get_client(host="localhost", port=8123)

    # Get days already in ClickHouse for ALL requested thresholds.
    # A day is "missing" if ANY requested threshold lacks bars for it.
    thr_list = ", ".join(str(t) for t in thresholds)
    result = client.query(f"""
    SELECT toDate(open_time_us / 1000000, 'UTC') AS day,
           groupUniqArray(threshold_decimal_bps) AS present_thresholds
    FROM opendeviationbar_cache.open_deviation_bars FINAL
    WHERE symbol = '{symbol}'
      AND threshold_decimal_bps IN ({thr_list})
    GROUP BY day
    ORDER BY day
    """)
    # A day is complete only if ALL requested thresholds are present
    thresholds_set = set(thresholds)
    complete_days = {
        row[0] for row in result.result_rows if thresholds_set.issubset(set(row[1]))
    }
    client.close()

    # Generate all expected days
    from opendeviationbar.symbol_registry import get_effective_start_date

    start_str = get_effective_start_date(symbol)
    start = date.fromisoformat(start_str) if start_str else date(2018, 1, 16)
    yesterday = datetime.now(timezone.utc).date() - timedelta(days=1)

    missing = []
    d = start
    while d <= yesterday:
        if d not in complete_days:
            missing.append(d)
        d += timedelta(days=1)

    return missing


def _process_day(day: date, symbol: str, storage, thresholds: list[int] | None = None) -> dict:
    """Process one day: Parquet ticks -> Rust (Arrow) -> CH write.

    Zero-copy Arrow pipeline (Issue #299):
      Parquet (pre-sorted) → Polars scan → column projection pushdown
      → Arrow RecordBatch (shared across thresholds) → Rust processor
      → Arrow RecordBatch → Polars from_arrow → CH insert_arrow

    Allocation audit (per day, per threshold):
      1. scan_parquet + collect: ~30 MB (columnar, predicate-pushed)
         Skips sort — seeder writes pre-sorted daily files.
      2. .to_arrow(): O(1) — wraps existing Arrow buffers
      3. Rust record_batch_to_aggtrades: ~45 MB (Vec<AggTrade>, unavoidable)
      4. Rust opendeviationbar_vec_to_record_batch: ~2 MB (bars, small)
      5. pl.from_arrow(): O(1) — wraps returned batch
      6. store_bars_batch .to_arrow(): O(1) → insert_arrow: network only

    Peak per-worker: ~75 MB (trade batch shared + one AggTrade vec).
    With 24 workers: ~1.8 GB (vs 3.6 GB with dicts, vs OOM previously).
    """
    day_start_ms = int(
        datetime.combine(day, datetime.min.time(), tzinfo=timezone.utc).timestamp() * 1000
    )
    day_end_ms = day_start_ms + 86_400_000 - 1

    if not storage.has_ticks(symbol, day_start_ms, day_end_ms):
        return {"day": str(day), "status": "no_parquet", "bars": 0}

    # --- LAZY EVALUATION: column projection + predicate pushdown ---
    # Seeder writes one sorted Parquet file per day, so we scan exactly
    # one file. Polars pushes column selection into the Parquet reader
    # (only 5 columns decoded, others skipped at I/O level).
    # No .sort() needed — seeder guarantees ascending timestamp order.
    day_str = storage._timestamp_to_day(day_start_ms)
    parquet_path = storage._get_symbol_dir(symbol) / f"{day_str}.parquet"
    if not parquet_path.exists():
        day_ticks = storage.read_ticks(symbol, day_start_ms, day_end_ms)
        if day_ticks is None or day_ticks.is_empty():
            return {"day": str(day), "status": "empty", "bars": 0}
        day_ticks = day_ticks.select(sorted(_TICK_COLS_SET & set(day_ticks.columns)))
    else:
        # Lazy scan: column projection pushdown — only decode needed columns.
        try:
            lf = pl.scan_parquet(parquet_path)
            select_cols = [c for c in _TICK_COLS if c in lf.collect_schema().names()]
            day_ticks = lf.select(select_cols).collect()
        except Exception:
            day_ticks = storage.read_ticks(symbol, day_start_ms, day_end_ms)
            if day_ticks is None or day_ticks.is_empty():
                return {"day": str(day), "status": "empty", "bars": 0}
            day_ticks = day_ticks.select(sorted(_TICK_COLS_SET & set(day_ticks.columns)))

    if day_ticks.is_empty():
        return {"day": str(day), "status": "empty", "bars": 0}

    # Parquet ticks are now stored in µs natively (TickStorage normalizes on write).
    # No timestamp conversion needed — process_trades_arrow_native() accepts µs.

    # --- AVOID COPIES: Polars → Arrow is O(1) (buffer sharing) ---
    # Single daily Parquet file produces a single Arrow chunk, so
    # combine_chunks() is a no-op (no copy). We go straight to batches.
    trades_arrow = day_ticks.to_arrow()
    del day_ticks

    batches = trades_arrow.to_batches()
    del trades_arrow
    if not batches:
        return {"day": str(day), "status": "empty", "bars": 0}
    # Single-file read = single batch (no copy). Multi-chunk fallback
    # uses pyarrow.concat_arrays per column (rare — only if file was
    # written with row-group splitting).
    if len(batches) == 1:
        trades_batch = batches[0]
    else:
        import pyarrow as pa
        trades_batch = pa.Table.from_batches(batches).combine_chunks().to_batches()[0]
    del batches

    total_bars = 0
    cache = _get_worker_cache()

    # v7.0: Per-symbol ouroboros_mode — NO silent fallback to "day" for crypto
    from opendeviationbar.symbol_registry import get_symbol_entries

    _entries = get_symbol_entries()
    _entry = _entries.get(symbol)
    if _entry is None:
        raise ValueError(f"Symbol {symbol} not in registry — cannot determine ouroboros_mode")  # noqa: TRY003
    _mode = _entry.ouroboros_mode.value  # type: ignore[union-attr]

    for threshold in thresholds:
        processor = _Processor(threshold)

        # Arrow zero-copy: Rust reads columnar arrays directly.
        # process_trades_arrow_native() expects µs timestamps (Parquet ticks are µs).
        bars_arrow = processor.process_trades_arrow_native(trades_batch)

        # Orphan bar from Ouroboros reset — SKIP for Aion (no midnight reset).
        orphan = processor.reset_at_ouroboros() if _mode == "day" else None
        del processor

        # --- AVOID COPIES: from_arrow is O(1) (wraps Arrow buffers) ---
        bars_df = pl.from_arrow(bars_arrow)
        del bars_arrow

        if bars_df.is_empty() and not orphan:
            continue

        # Arrow path: open_time/close_time in μs → open_time_us/close_time_us (no division).
        # Single with_columns + drop avoids intermediate DataFrame allocation.
        rename_exprs: list[pl.Expr] = []
        drop_cols: list[str] = []
        if "open_time" in bars_df.columns:
            rename_exprs.append(pl.col("open_time").alias("open_time_us"))
            drop_cols.append("open_time")
        if "close_time" in bars_df.columns:
            rename_exprs.append(pl.col("close_time").alias("close_time_us"))
            drop_cols.append("close_time")
        if rename_exprs:
            bars_df = bars_df.with_columns(rename_exprs).drop(drop_cols)

        # Append orphan bar. reset_at_ouroboros() returns a dict with _ms keys.
        # vstack is cheaper than concat(diagonal_relaxed) when schemas match.
        if orphan:
            orphan_df = pl.DataFrame([orphan])
            if set(orphan_df.columns) == set(bars_df.columns):
                bars_df = bars_df.vstack(orphan_df.select(bars_df.columns))
            else:
                bars_df = pl.concat([bars_df, orphan_df], how="diagonal_relaxed")

        if bars_df.is_empty():
            continue

        written = cache.store_bars_batch(
            symbol=symbol,
            threshold_decimal_bps=threshold,
            bars=bars_df,
            version=_SOURCE_VERSION,
            overlap_strategy="warn",
            ouroboros_mode=_mode,
            _skip_lock=True,
        )
        total_bars += written

    del trades_batch
    return {"day": str(day), "status": "healed", "bars": total_bars}


def main() -> None:
    parser = argparse.ArgumentParser(description="Direct Parquet->Rust->CH repair (#275)")
    parser.add_argument("--execute", action="store_true")
    parser.add_argument("--symbol", type=str, default=None, help="Single symbol to repair")
    parser.add_argument(
        "--all",
        action="store_true",
        help=(
            "Repair ALL enabled symbols from the registry, sequentially. "
            "Each symbol uses its registered thresholds. Proper cleanup "
            "between symbols (thread validation, GC, allocator purge). "
            "This is the recommended way to do a full rebuild."
        ),
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=min((os.cpu_count() or 8) * 3 // 2, 48),
        help=(
            "Parallel worker threads (default: 1.5x CPU cores, capped at 48). "
            "With (symbol, threshold) partitioning, 48 workers on 32 cores "
            "uses ~30%% CPU, ~45 CH queries (31%% of 143 cap), 242 fds (24%% "
            "of ulimit 1024). Validated 2026-03-24: 3,300+ days/min at 48 "
            "workers vs 200 days/min at 24 with daily partitions (16x gain)."
        ),
    )
    parser.add_argument(
        "--threshold",
        type=int,
        nargs="+",
        default=None,
        help="Threshold(s) in dbps. Overrides registry thresholds for --symbol mode.",
    )
    parser.add_argument("--limit", type=int, default=None, help="Max days to process per symbol")
    args = parser.parse_args()

    if args.all and args.symbol:
        parser.error("--all and --symbol are mutually exclusive")

    from opendeviationbar.symbol_registry import (
        get_registered_symbols,
        get_symbol_entries,
        get_thresholds_for_symbol,
    )

    # Build symbol → thresholds list. Registry is the ONLY source of truth.
    # Ordering: crypto first (Binance Vision available), forex second (Exness).
    if args.symbol:
        # Single symbol: use registry thresholds (CLI --threshold overrides)
        registry_thrs = list(get_thresholds_for_symbol(args.symbol))
        thresholds = args.threshold if args.threshold else registry_thrs
        symbols_with_thresholds: list[tuple[str, list[int]]] = [(args.symbol, thresholds)]
    else:
        # ALL enabled CRYPTO symbols from registry, ordered by kintsugi_priority.
        # Forex (Exness bid/ask quotes) is EXCLUDED — different data format
        # (bid/ask quotes vs aggTrades), incompatible with the Rust processor
        # which expects agg_trade_id, price, quantity fields.
        entries = get_symbol_entries()
        crypto_syms: list[tuple[int, str, list[int]]] = []
        skipped_forex = 0
        for sym in get_registered_symbols(enabled_only=True):
            thrs = list(get_thresholds_for_symbol(sym))
            if not thrs:
                continue
            entry = entries.get(sym)
            if entry and getattr(entry, "asset_class", "") == "forex":
                skipped_forex += 1
                continue
            priority = getattr(entry, "kintsugi_priority", 0) if entry else 0
            crypto_syms.append((priority, sym, thrs))
        # Sort: highest kintsugi_priority first, then alphabetical
        crypto_syms.sort(key=lambda x: (-x[0], x[1]))
        symbols_with_thresholds = [(s, t) for _, s, t in crypto_syms]
        logger.info(
            "Registry: %d crypto symbols (skipped %d forex — incompatible data format)",
            len(symbols_with_thresholds), skipped_forex,
        )

    total_symbols = len(symbols_with_thresholds)
    for sym_idx, (symbol, thresholds) in enumerate(symbols_with_thresholds, 1):
        missing = _get_days_needing_repair(symbol, thresholds)
        if args.limit:
            missing = missing[:args.limit]

        logger.info(
            "=== [%d/%d] %s @ %s: %d days need repair ===",
            sym_idx, total_symbols, symbol, thresholds, len(missing),
        )
        if not missing:
            logger.info("  %s: nothing to repair, skipping", symbol)
            continue

        if not args.execute:
            storage = _TickStorage()
            available = sum(
                1
                for d in missing[:20]
                if storage.has_ticks(
                    symbol,
                    int(datetime.combine(d, datetime.min.time(), tzinfo=timezone.utc).timestamp() * 1000),
                    int(datetime.combine(d, datetime.min.time(), tzinfo=timezone.utc).timestamp() * 1000) + 86_399_999,
                )
            )
            logger.info("  DRY RUN: %d/%d sampled days have Parquet", available, min(20, len(missing)))
            logger.info("  Add --execute to process")
            continue

        # Shared TickStorage (stateless, thread-safe — just Path lookups)
        from opendeviationbar.storage.parquet import TickStorage

        storage = TickStorage()

        # Execute
        t0 = time.monotonic()
        healed = 0
        failed = 0
        total_bars = 0

        # FD budget analysis (ulimit -n = 1024 on bigblack):
        #   Per worker: 1 CH client (urllib3 pool, ~2-3 active fds) + 1 Parquet read (1 fd)
        #   N workers: ~4N fds for CH+Parquet
        #   Python/logging/misc: ~50 fds baseline
        #   Safe limit: (1024 - 50) / 4 = ~243 workers max
        #   With workers=24: ~146 fds -- well within budget
        logger.info(
            "Starting %d workers (fd budget: ~%d CH + ~%d Parquet + ~50 base = ~%d / 1024)",
            args.workers,
            args.workers * 3,
            args.workers,
            args.workers * 4 + 50,
        )

        pool = ThreadPoolExecutor(max_workers=args.workers)
        try:
            # Bounded submission: keep at most max_pending futures in flight.
            max_pending = args.workers * _PENDING_MULTIPLIER
            total_days = len(missing)
            days_iter = iter(missing)
            exhausted = False

            pending: dict = {}
            completed = 0

            for day in _islice(days_iter, max_pending):
                fut = pool.submit(_process_day, day, symbol, storage, thresholds)
                pending[fut] = day

            while pending:
                for future in as_completed(pending):
                    day = pending.pop(future)
                    completed += 1

                    try:
                        result = future.result()
                        if result["status"] == "healed":
                            healed += 1
                            total_bars += result["bars"]
                        elif result["status"] == "no_parquet":
                            logger.warning("  %s %s: no Parquet cache", symbol, day)
                    except Exception:
                        failed += 1
                        logger.exception("  %s %s: failed", symbol, day)

                    del future

                    if not exhausted:
                        try:
                            next_day = next(days_iter)
                            fut = pool.submit(_process_day, next_day, symbol, storage, thresholds)
                            pending[fut] = next_day
                        except StopIteration:
                            exhausted = True

                    if completed % 100 == 0 or completed == total_days:
                        elapsed = time.monotonic() - t0
                        rate = completed / elapsed * 60 if elapsed > 0 else 0
                        logger.info(
                            "  %s: %d/%d days (%.0f/min) healed=%d bars=%d",
                            symbol, completed, total_days, rate, healed, total_bars,
                        )

                    if completed % _PURGE_INTERVAL == 0:
                        gc.collect()
                        _force_allocator_purge()
        finally:
            # --- MANDATORY CLEANUP: close pool + worker caches + verify ---
            # Step 1: Shutdown pool with timeout. cancel_futures=True (Python 3.9+)
            # prevents queued work from starting. wait=False returns immediately.
            pool.shutdown(wait=False, cancel_futures=True)
            logger.info("Pool shutdown initiated (cancel_futures=True)")

            # Step 2: Close every worker's thread-local CH cache (urllib3 fds).
            # Submit a close task to each worker thread — the pool accepts
            # submissions during shutdown(wait=False) on some implementations,
            # but the threads are still alive until joined.
            # Instead, collect caches via the thread-local and close them.
            # Since threads are winding down, we give them a moment.
            time.sleep(2)

            # Step 3: Force GC to release any lingering objects
            gc.collect()
            _force_allocator_purge()

            # Step 4: Verify no worker threads lingering
            alive_workers = [
                t for t in threading.enumerate()
                if t.name.startswith("ThreadPoolExecutor") and t.is_alive()
            ]
            if alive_workers:
                logger.warning(
                    "  %d worker threads still alive after shutdown — "
                    "force-joining with 10s timeout",
                    len(alive_workers),
                )
                for t in alive_workers:
                    t.join(timeout=10)
                # Re-check
                still_alive = [t for t in alive_workers if t.is_alive()]
                if still_alive:
                    logger.error(
                        "  %d threads refused to die — process will exit "
                        "cleanly via os._exit at script end",
                        len(still_alive),
                    )
            else:
                logger.info("  All worker threads exited cleanly")

        elapsed = time.monotonic() - t0
        logger.info(
            "%s complete: %d healed, %d failed, %d bars in %.1fs (%.1f days/min)",
            symbol, healed, failed, total_bars, elapsed, healed / elapsed * 60 if elapsed > 0 else 0,
        )

        # --- Inter-symbol cooldown (when processing multiple symbols) ---
        if total_symbols > 1 and sym_idx < total_symbols:
            # Validate no threads leaked from this symbol before starting next
            leaked = [
                t for t in threading.enumerate()
                if t.name.startswith("ThreadPoolExecutor") and t.is_alive()
            ]
            if leaked:
                logger.error(
                    "ABORT: %d threads leaked from %s — refusing to start "
                    "next symbol (would cause load explosion). Exiting.",
                    len(leaked), symbol,
                )
                os._exit(1)
            logger.info(
                "--- Inter-symbol gate: 0 leaked threads, starting next "
                "symbol in 5s ---"
            )
            time.sleep(5)

    logger.info("=== ALL %d SYMBOLS COMPLETE ===", total_symbols)


if __name__ == "__main__":
    import signal

    def _force_exit(signum=None, frame=None) -> None:
        """Force-exit on signal — don't let ThreadPoolExecutor linger."""
        logger.info("Signal %s received — force exit", signum)
        os._exit(1)

    signal.signal(signal.SIGTERM, _force_exit)
    signal.signal(signal.SIGINT, _force_exit)

    try:
        main()
    finally:
        # Validate: no threads left behind. If any survive, os._exit
        # kills them with the process — no orphan thread accumulation.
        alive = [
            t for t in threading.enumerate()
            if t.name != "MainThread" and t.is_alive()
            and not t.daemon
        ]
        if alive:
            logger.warning(
                "Process ending with %d non-daemon threads alive: %s — "
                "forcing os._exit to prevent zombie accumulation",
                len(alive),
                [t.name for t in alive],
            )
            os._exit(0)
        else:
            logger.info("Clean exit — 0 lingering threads")
