#!/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.
# 50 is a good balance: not so frequent it hurts throughput, not so rare
# that arena fragmentation accumulates 20+ GB before the first purge.
_PURGE_INTERVAL = 50

THRESHOLDS = [250, 500, 750, 1000]

# 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) -> 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 (any threshold counts)
    result = client.query(f"""
    SELECT DISTINCT toDate(open_time_ms / 1000, 'UTC') AS day
    FROM opendeviationbar_cache.open_deviation_bars FINAL
    WHERE symbol = '{symbol}'
    ORDER BY day
    """)
    existing_days = {row[0] for row in result.result_rows}
    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 existing_days:
            missing.append(d)
        d += timedelta(days=1)

    return missing


def _process_day(day: date, symbol: str, storage) -> dict:
    """Process one day: Parquet ticks -> 4 threshold processors -> CH write.

    Same approach as kintsugi._repair_under_lock() Parquet path, but processes
    all 4 thresholds from a single tick read (4x less I/O).

    Uses thread-local OpenDeviationBarCache to avoid fd exhaustion.
    Uses .to_dicts() for Polars->Rust interop (Arrow zero-copy deferred --
    process_trades_arrow() output schema doesn't yet match store_bars_batch()).
    Uses overlap_strategy="warn" (plain INSERT) to avoid 300s DELETE scan
    timeout with "replace" on 1,658+ active parts.
    """
    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}

    # Read ticks ONCE for all thresholds
    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}

    # Select only trade columns needed by the processor
    needed = {"agg_trade_id", "price", "quantity", "timestamp", "is_buyer_maker"}
    available = set(day_ticks.columns)
    day_ticks = day_ticks.select(sorted(needed & available))

    total_bars = 0
    source_version = f"{__version__}+repair-parquet"

    # Convert to dicts ONCE, reuse across all 4 thresholds.
    # Dict path: ~150 MB per day for BTCUSDT (760K trades * ~200 bytes/dict).
    # Bounded futures + allocator purge keep aggregate RSS manageable (~13 GB).
    trades = day_ticks.to_dicts()
    del day_ticks  # Free Polars DataFrame immediately

    # Reuse thread-local cache for all 4 thresholds (was: new cache per threshold)
    cache = _get_worker_cache()

    for threshold in THRESHOLDS:
        processor = Processor(threshold)
        bars = processor.process_trades(trades)

        orphan = processor.reset_at_ouroboros()
        if orphan:
            bars.append(orphan)
        del processor

        if not bars:
            continue

        bars_df = pl.DataFrame(bars)
        del bars

        if bars_df.is_empty():
            del bars_df
            continue

        written = cache.store_bars_batch(
            symbol=symbol,
            threshold_decimal_bps=threshold,
            bars=bars_df,
            version=source_version,
            overlap_strategy="warn",  # Plain INSERT — no DELETE scan (remaining days have 0 existing bars)
            _skip_lock=True,
        )
        total_bars += written
        del bars_df

    del trades
    gc.collect()

    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("--limit", type=int, default=None, help="Max days to process")
    args = parser.parse_args()

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

    for symbol in symbols:
        missing = _get_days_needing_repair(symbol)
        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)
                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)
                            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()
