#!/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

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

# Bounded submission: keep at most 2*workers futures pending.
# With 24 workers this means 48 futures max, not 2,000+.
_PENDING_MULTIPLIER = 2

# 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]

# 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


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

    Each worker cycles through Polars DataFrames + Python dicts per day.
    Without periodic purge, malloc arena fragmentation can inflate RSS
    well beyond live object size. This function compacts the heap, returning
    unused pages to the OS.

    On Linux with mimalloc (LD_PRELOAD), calls mi_collect(true) for
    immediate page decommit. Falls back to glibc malloc_trim(0).
    No-op on macOS (development machines don't hit this).
    """
    import ctypes
    import sys

    if sys.platform != "linux":
        return

    # Try mimalloc first (loaded via LD_PRELOAD in systemd service)
    try:
        libmi = ctypes.CDLL("libmimalloc.so.2")
        libmi.mi_collect(ctypes.c_bool(True))  # force=true: immediate purge
    except OSError:
        pass
    else:
        return

    # Fallback: glibc malloc_trim (returns free pages at top of heap)
    try:
        libc = ctypes.CDLL("libc.so.6")
        libc.malloc_trim(0)
    except OSError:
        pass


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_ms / 1000, '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).
    """
    import polars as pl
    from opendeviationbar import __version__
    from opendeviationbar._core import PyOpenDeviationBarProcessor as Processor

    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.
    tick_cols = ["agg_trade_id", "price", "quantity", "timestamp", "is_buyer_maker"]
    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():
        # Fallback to the full read_ticks path if path resolution fails
        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}
        available = set(day_ticks.columns)
        day_ticks = day_ticks.select(sorted({"agg_trade_id", "price", "quantity", "timestamp", "is_buyer_maker"} & available))
    else:
        # Direct scan: column projection pushdown avoids decoding unused cols.
        # Single daily file = single Arrow chunk = no combine_chunks() needed.
        try:
            lf = pl.scan_parquet(parquet_path)
            schema_names = lf.collect_schema().names()
            select_cols = [c for c in tick_cols if c in 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}
            available = set(day_ticks.columns)
            day_ticks = day_ticks.select(sorted({"agg_trade_id", "price", "quantity", "timestamp", "is_buyer_maker"} & available))

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

    # --- 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
    source_version = f"{__version__}+repair-parquet"
    cache = _get_worker_cache()

    for threshold in (thresholds or THRESHOLDS_DEFAULT):
        processor = Processor(threshold)

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

        # Orphan bar from Ouroboros reset (single Python dict, negligible).
        orphan = processor.reset_at_ouroboros()
        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_ms/close_time_ms in ms.
        # 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") // 1000).alias("open_time_ms"))
            drop_cols.append("open_time")
        if "close_time" in bars_df.columns:
            rename_exprs.append((pl.col("close_time") // 1000).alias("close_time_ms"))
            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",
            _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)
    parser.add_argument(
        "--workers",
        type=int,
        default=min(os.cpu_count() or 8, 24),
        help=(
            "Parallel worker threads (default: min(nproc, 24)). "
            "24 is safe on bigblack (32 cores): thread-local CH cache keeps fds "
            "at ~4N+50 (well under ulimit -n 1024), and each worker is mostly "
            "I/O-bound (Parquet read + CH write), so 24 on 32 cores is fine. "
            "Do NOT exceed 24 — ClickHouse query scheduler saturates at ~143 "
            "concurrent queries (32 workers = 0 completions)."
        ),
    )
    parser.add_argument(
        "--threshold",
        type=int,
        nargs="+",
        default=None,
        help="Threshold(s) in dbps (default: 250 500 750). Example: --threshold 100",
    )
    parser.add_argument("--limit", type=int, default=None, help="Max days to process")
    args = parser.parse_args()

    thresholds = args.threshold if args.threshold else THRESHOLDS_DEFAULT
    symbols = [args.symbol] if args.symbol else ["BTCUSDT", "ETHUSDT"]

    for symbol in symbols:
        missing = _get_days_needing_repair(symbol, thresholds)
        if args.limit:
            missing = missing[:args.limit]

        logger.info("=== %s: %d days need repair ===", symbol, len(missing))
        if not missing:
            continue

        if not args.execute:
            from opendeviationbar.storage.parquet import TickStorage

            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,
        )

        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            # Bounded submission: keep at most max_pending futures in flight.
            # Old pattern submitted all 2,000+ days at once, keeping every
            # Future (and its result) alive in a dict until the loop ended.
            # New pattern: sliding window of pending futures. Completed
            # futures are discarded immediately, allowing GC of their
            # result data and the closure args.
            max_pending = args.workers * _PENDING_MULTIPLIER
            total_days = len(missing)
            days_iter = iter(missing)
            exhausted = False  # True once days_iter is consumed

            # {Future -> date} for pending futures only
            pending: dict = {}
            completed = 0

            # Seed the pool with initial batch
            for day in _islice(days_iter, max_pending):
                fut = pool.submit(_process_day, day, symbol, storage, thresholds)
                pending[fut] = day

            # Drain all futures. as_completed() takes a snapshot internally,
            # so it is safe that we mutate `pending` during iteration. We
            # let it run to exhaustion (processes all currently-pending
            # futures), then loop back to check if new ones were added.
            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)

                    # Explicitly drop the future to release result + closure refs
                    del future

                    # Submit next day to keep the pool saturated
                    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

                    # Progress logging
                    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,
                        )

                    # Periodic allocator purge to prevent RSS bloat from
                    # malloc arena fragmentation. gc.collect() frees Python
                    # objects but glibc/mimalloc hold the pages. This forces
                    # page return to OS.
                    if completed % _PURGE_INTERVAL == 0:
                        gc.collect()
                        _force_allocator_purge()
                        logger.debug(
                            "  Allocator purge after %d completions", completed,
                        )

        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,
        )


if __name__ == "__main__":
    main()
