#!/usr/bin/env python3
# FILE-SIZE-OK: Cohesive repair script — Aion processors, CH write, checkpoint save
"""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, and
writes bars to ClickHouse.

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

Aion-only: maintains continuous processors across day boundaries (no
midnight reset, no orphan bars). Parallelism is at the SYMBOL level.

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:
  - Allocator purge: _force_allocator_purge() periodically 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.

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 json
import logging
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, datetime, timedelta, timezone
from pathlib import Path

# 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

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

# 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)
    # Aion: include today (sidecar resumes from repair's last TID).
    end_date = datetime.now(timezone.utc).date()

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

    return missing


def _read_day_trades(day: date, symbol: str, storage) -> object | None:
    """Read one day's Parquet trades as an Arrow RecordBatch. Returns None if unavailable."""
    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 None

    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 None
        day_ticks = day_ticks.select(sorted(_TICK_COLS_SET & set(day_ticks.columns)))
    else:
        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 (OSError, RuntimeError):
            day_ticks = storage.read_ticks(symbol, day_start_ms, day_end_ms)
            if day_ticks is None or day_ticks.is_empty():
                return None
            day_ticks = day_ticks.select(sorted(_TICK_COLS_SET & set(day_ticks.columns)))

    if day_ticks.is_empty():
        return None

    trades_arrow = day_ticks.to_arrow()
    del day_ticks
    batches = trades_arrow.to_batches()
    del trades_arrow
    if not batches:
        return None
    if len(batches) == 1:
        return batches[0]
    import pyarrow as pa

    return pa.Table.from_batches(batches).combine_chunks().to_batches()[0]


def _process_symbol_aion(
    symbol: str,
    days: list[date],
    storage,
    thresholds: list[int],
    *,
    checkpoints_only: bool = False,
) -> dict:
    """Process ALL days for a symbol through continuous processors.

    Maintains ONE processor per threshold across ALL days. The processor
    state carries across day boundaries — the forming bar at midnight is
    NOT emitted as an orphan but continues into the next day.

    Result: ZERO day-boundary gaps. Every trade belongs to exactly one bar.

    Days MUST be sorted chronologically (earliest first). This is sequential
    per symbol — parallelism is at the SYMBOL level (different symbols can
    run concurrently), not the DAY level.
    """
    cache = _get_worker_cache()

    from opendeviationbar.symbol_registry import get_symbol_entries

    _entry = get_symbol_entries().get(symbol)
    if _entry is None:
        msg = f"Symbol {symbol} not in registry"
        raise ValueError(msg)
    _mode = _entry.ouroboros_mode.value

    total_bars = 0
    total_days = len(days)

    # Sort chronologically — essential for continuous processor state
    days_sorted = sorted(days)

    # ONE processor per threshold, persistent across ALL days
    processors = {thr: _Processor(thr) for thr in thresholds}

    t0 = time.monotonic()
    for day_idx, day in enumerate(days_sorted):
        trades_batch = _read_day_trades(day, symbol, storage)
        if trades_batch is None:
            continue

        # Feed the SAME day's trades to all threshold processors
        for thr in thresholds:
            processor = processors[thr]
            bars_arrow = processor.process_trades_arrow_native(trades_batch)

            bars_df = pl.from_arrow(bars_arrow)
            del bars_arrow

            if bars_df.is_empty():
                continue

            # Filter out Rust processor initialization bars (first_agg_trade_id=0).
            # The continuous processor emits a spurious bar before real trades arrive.
            if "first_agg_trade_id" in bars_df.columns:
                bars_df = bars_df.filter(pl.col("first_agg_trade_id") > 0)
                if bars_df.is_empty():
                    continue

            if not checkpoints_only:
                # Rename columns (Arrow uses open_time/close_time, CH uses _us suffix)
                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)

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

        del trades_batch

        # Progress logging every 100 days
        if (day_idx + 1) % 100 == 0 or (day_idx + 1) == total_days:
            elapsed = time.monotonic() - t0
            rate = (day_idx + 1) / elapsed * 60 if elapsed > 0 else 0
            logger.info(
                "  %s [aion]: %d/%d days (%.0f/min) bars=%d",
                symbol, day_idx + 1, total_days, rate, total_bars,
            )

        # Periodic GC + allocator purge
        if (day_idx + 1) % _PURGE_INTERVAL == 0:
            gc.collect()
            _force_allocator_purge()

    # Save processor checkpoints so sidecar can continue with exact same
    # forming bar state — zero junction gaps on restart (Issue #317).
    try:
        from platformdirs import user_cache_dir

        ckpt_dir = Path(user_cache_dir("opendeviationbar")) / "repair_checkpoints"
        ckpt_dir.mkdir(parents=True, exist_ok=True)
        for thr, processor in processors.items():
            ckpt = processor.create_checkpoint(symbol)
            ckpt_path = ckpt_dir / f"{symbol}_{thr}.json"
            ckpt_path.write_text(json.dumps(ckpt, default=str))
        logger.info(
            "%s [aion]: saved %d processor checkpoints to %s",
            symbol, len(processors), ckpt_dir,
        )
    except Exception:
        logger.warning(
            "%s: failed to save processor checkpoints (non-fatal)", symbol,
            exc_info=True,
        )

    # Cleanup processors — NO orphan emission for Aion
    del processors

    elapsed = time.monotonic() - t0
    logger.info(
        "%s [aion] complete: %d bars from %d days in %.1fs (%.1f days/min)",
        symbol, total_bars, total_days, elapsed,
        total_days / elapsed * 60 if elapsed > 0 else 0,
    )
    return {"symbol": symbol, "status": "healed", "bars": total_bars, "days": total_days}


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(
        "--checkpoints-only",
        action="store_true",
        help=(
            "Process all trades through Aion processors but SKIP writing bars "
            "to CH. Only saves processor checkpoints for sidecar handoff. "
            "Much faster than full repair — no CH writes, no cleanup needed."
        ),
    )
    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)

    # Parallel across symbols, sequential within each.
    # Each symbol gets ONE thread that runs _process_symbol_aion() sequentially
    # through all its days. This gives N-way parallelism (N=16 crypto symbols)
    # while maintaining continuous processor state per symbol.
    from opendeviationbar.storage.parquet import TickStorage

    if args.execute:
        storage = TickStorage()
        # Build work items: (symbol, days, thresholds)
        work = []
        for symbol, thresholds in symbols_with_thresholds:
            missing = _get_days_needing_repair(symbol, thresholds)
            if args.limit:
                missing = missing[:args.limit]
            if missing:
                work.append((symbol, missing, thresholds))
                logger.info("  %s: %d days to rebuild (Aion continuous)", symbol, len(missing))
            else:
                logger.info("  %s: nothing to repair", symbol)

        if work:
            # Parallel: one thread per symbol (max = number of symbols)
            max_parallel = min(len(work), args.workers, 16)
            logger.info(
                "=== Parallel rebuild: %d symbols, %d threads ===",
                len(work), max_parallel,
            )
            t0 = time.monotonic()

            with ThreadPoolExecutor(max_workers=max_parallel) as pool:
                futures = {}
                for symbol, missing, thresholds in work:
                    fut = pool.submit(
                        _process_symbol_aion, symbol, missing, storage, thresholds,
                        checkpoints_only=args.checkpoints_only,
                    )
                    futures[fut] = symbol

                for fut in as_completed(futures):
                    sym = futures[fut]
                    try:
                        result = fut.result()
                        logger.info(
                            "=== %s complete: %d bars from %d days ===",
                            sym, result["bars"], result["days"],
                        )
                    except (OSError, RuntimeError, ValueError):
                        logger.exception("  %s: FAILED", sym)

            elapsed = time.monotonic() - t0
            logger.info(
                "=== ALL %d symbols complete in %.1fs ===",
                len(work), elapsed,
            )
            gc.collect()
            _force_allocator_purge()

    else:
        storage = TickStorage()
        for symbol, thresholds in symbols_with_thresholds:
            missing = _get_days_needing_repair(symbol, thresholds)
            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(
                "  %s: DRY RUN — %d days need repair, %d/%d have Parquet",
                symbol, len(missing), available, min(20, len(missing)),
            )
        logger.info("  Add --execute to process")

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