#!/usr/bin/env python3
"""Fleet Recompute: Aion Migration for ClickHouse (v7.0, one-time).

Coordinated stop-drop-wipe-recompute script for migrating all crypto symbols
from Day-mode to Aion (continuous) mode on bigblack. This is a one-time
migration tool -- not a permanent operational utility.

Sequence:
    1. --preflight  : Discover partitions, show what will be dropped
    2. --execute    : DROP PARTITION + wipe telemetry (requires manual confirmation)
    3. --verify     : Post-recompute Stathera audit (0-gap 0-overlap)

The script connects to ClickHouse directly (local or via SSH tunnel).
systemctl commands are displayed as instructions -- the user must run them
manually on bigblack.

Usage:
    python scripts/fleet_recompute_aion.py --preflight
    python scripts/fleet_recompute_aion.py --execute
    python scripts/fleet_recompute_aion.py --execute --dry-run
    python scripts/fleet_recompute_aion.py --verify

Environment:
    OPENDEVIATIONBAR_MODE      Connection mode: local (bigblack), remote (SSH tunnel)
    OPENDEVIATIONBAR_CH_HOSTS  SSH alias (default: bigblack)
    CLICKHOUSE_HOST            Direct host override (default: localhost)
    CLICKHOUSE_PORT            Port override (default: 8123)
"""

from __future__ import annotations

import argparse
import contextlib
import os
import pathlib
import socket
import subprocess
import sys
import time
from datetime import UTC, datetime
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import clickhouse_connect

# =============================================================================
# ANSI Colors
# =============================================================================

GREEN = "\033[92m"
RED = "\033[91m"
YELLOW = "\033[93m"
CYAN = "\033[96m"
BOLD = "\033[1m"
RESET = "\033[0m"


def _ok(msg: str) -> None:
    print(f"  {GREEN}[OK]{RESET} {msg}")


def _warn(msg: str) -> None:
    print(f"  {YELLOW}[WARN]{RESET} {msg}")


def _err(msg: str) -> None:
    print(f"  {RED}[FAIL]{RESET} {msg}")


def _info(msg: str) -> None:
    print(f"  {CYAN}[INFO]{RESET} {msg}")


def _header(msg: str) -> None:
    print(f"\n{BOLD}{msg}{RESET}")
    print("=" * len(msg))


# =============================================================================
# Connection (reused from detect_gaps.py pattern)
# =============================================================================


def _find_free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("localhost", 0))
        return s.getsockname()[1]


def _is_port_open(host: str, port: int, timeout: float = 1.0) -> bool:
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(timeout)
            return s.connect_ex((host, port)) == 0
    except OSError:
        return False


class SSHTunnel:
    """Minimal SSH tunnel manager for ClickHouse access."""

    def __init__(self, ssh_alias: str, remote_port: int = 8123) -> None:
        self.ssh_alias = ssh_alias
        self.remote_port = remote_port
        self.local_port: int | None = None
        self._process: subprocess.Popen[bytes] | None = None

    def start(self, timeout: float = 10.0) -> int:
        self.local_port = _find_free_port()
        self._process = subprocess.Popen(
            [
                "ssh",
                "-N",
                "-o", "ExitOnForwardFailure=yes",
                "-o", "ServerAliveInterval=30",
                "-o", "ServerAliveCountMax=3",
                "-L", f"{self.local_port}:localhost:{self.remote_port}",
                self.ssh_alias,
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
        )
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            if _is_port_open("localhost", self.local_port):
                return self.local_port
            exit_code = self._process.poll()
            if exit_code is not None:
                if exit_code == 0:
                    time.sleep(0.3)
                    if _is_port_open("localhost", self.local_port):
                        self._process = None
                        return self.local_port
                stderr = ""
                if self._process.stderr:
                    stderr = self._process.stderr.read().decode()
                self._process = None
                msg = f"SSH tunnel to {self.ssh_alias} failed (exit={exit_code}): {stderr}"
                raise RuntimeError(msg)
            time.sleep(0.2)
        self.stop()
        msg = f"SSH tunnel to {self.ssh_alias} timed out after {timeout}s"
        raise RuntimeError(msg)

    def stop(self) -> None:
        if self._process is not None:
            try:
                self._process.terminate()
                self._process.wait(timeout=5)
            except (subprocess.TimeoutExpired, OSError):
                try:
                    self._process.kill()
                    self._process.wait(timeout=2)
                except (subprocess.TimeoutExpired, OSError):
                    pass
            finally:
                for stream in (
                    self._process.stdin,
                    self._process.stdout,
                    self._process.stderr,
                ):
                    if stream is not None:
                        with contextlib.suppress(OSError):
                            stream.close()
                self._process = None


def _get_client(
    host: str = "localhost", port: int = 8123,
) -> clickhouse_connect.driver.Client:
    import clickhouse_connect

    client = clickhouse_connect.get_client(host=host, port=port)
    client.command("SELECT 1")
    return client


def connect() -> tuple[clickhouse_connect.driver.Client, SSHTunnel | None]:
    """Connect to ClickHouse based on OPENDEVIATIONBAR_MODE."""
    mode = os.getenv("OPENDEVIATIONBAR_MODE", "remote").lower()
    ch_host = os.getenv("CLICKHOUSE_HOST", "localhost")
    ch_port = int(os.getenv("CLICKHOUSE_PORT", "8123"))

    if mode == "local":
        return _get_client(ch_host, ch_port), None

    if mode == "remote":
        ssh_alias = (
            os.getenv("OPENDEVIATIONBAR_CH_HOSTS", "bigblack").split(",")[0].strip()
        )
        if _is_port_open("localhost", 18123, timeout=0.5):
            return _get_client("localhost", 18123), None
        try:
            result = subprocess.run(
                ["ssh", "-G", ssh_alias],
                capture_output=True, text=True, timeout=2, check=False,
            )
            for line in result.stdout.splitlines():
                if line.startswith("hostname "):
                    resolved_ip = line.split()[1]
                    if _is_port_open(resolved_ip, 8123, timeout=2.0):
                        return _get_client(resolved_ip, 8123), None
        except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
            pass
        tunnel = SSHTunnel(ssh_alias)
        local_port = tunnel.start()
        return _get_client("localhost", local_port), tunnel

    # auto mode
    if _is_port_open(ch_host, ch_port, timeout=1.0):
        return _get_client(ch_host, ch_port), None
    ssh_alias = os.getenv("OPENDEVIATIONBAR_CH_HOSTS", "bigblack").split(",")[0].strip()
    if ssh_alias:
        tunnel = SSHTunnel(ssh_alias)
        local_port = tunnel.start()
        return _get_client("localhost", local_port), tunnel
    msg = "Cannot connect to ClickHouse in auto mode"
    raise RuntimeError(msg)


# =============================================================================
# Symbol Discovery
# =============================================================================

def _get_crypto_symbols() -> list[str]:
    """Get all registered crypto symbols from the symbol registry."""
    from opendeviationbar.symbol_registry import get_registered_symbols

    return list(get_registered_symbols(asset_class="crypto"))


def _discover_crypto_partitions(
    client: clickhouse_connect.driver.Client,
    crypto_symbols: list[str],
) -> list[tuple[str, int]]:
    """Query system.parts to discover actual crypto partitions on disk.

    Returns list of (symbol, threshold_decimal_bps) tuples for partitions
    that actually exist in ClickHouse (not just theoretically expected).
    """
    result = client.query(
        """
        SELECT DISTINCT partition
        FROM system.parts
        WHERE table = 'open_deviation_bars'
          AND database = 'opendeviationbar_cache'
          AND active = 1
        ORDER BY partition
        """
    )

    partitions: list[tuple[str, int]] = []
    crypto_set = set(crypto_symbols)

    for row in result.result_rows:
        partition_str = row[0]
        # Partition format: ('SYMBOL', threshold) -> parse it
        # ClickHouse partition IDs look like: ('BTCUSDT', 250)
        # But system.parts stores it as a string representation
        try:
            # Try parsing tuple format: ('SYMBOL', 250)
            cleaned = partition_str.strip("()")
            parts = cleaned.split(",")
            if len(parts) == 2:
                symbol = parts[0].strip().strip("'\"")
                threshold = int(parts[1].strip())
                if symbol in crypto_set:
                    partitions.append((symbol, threshold))
        except (ValueError, IndexError):
            continue

    return sorted(partitions)


# =============================================================================
# Phase 1: Preflight
# =============================================================================

def run_preflight(client: clickhouse_connect.driver.Client) -> int:
    """Preflight check: discover partitions and show summary."""
    _header("Phase 1: Preflight Check")

    # 1. Verify CH connectivity
    try:
        client.command("SELECT 1")
        _ok("ClickHouse connectivity verified")
    except Exception as e:
        _err(f"ClickHouse connectivity failed: {e}")
        return 2

    # 2. Get crypto symbols from registry
    crypto_symbols = _get_crypto_symbols()
    _info(f"Registry crypto symbols: {len(crypto_symbols)}")
    for sym in crypto_symbols:
        print(f"    {sym}")

    # 3. Discover actual partitions
    partitions = _discover_crypto_partitions(client, crypto_symbols)
    _info(f"Actual crypto partitions on disk: {len(partitions)}")

    if not partitions:
        _warn("No crypto partitions found -- nothing to drop")
        return 0

    # Group by symbol for display
    by_symbol: dict[str, list[int]] = {}
    for sym, thr in partitions:
        by_symbol.setdefault(sym, []).append(thr)

    print(f"\n  {'Symbol':<16} {'Thresholds'}")
    print(f"  {'------':<16} {'----------'}")
    for sym in sorted(by_symbol):
        thrs = ", ".join(str(t) for t in sorted(by_symbol[sym]))
        print(f"  {sym:<16} {thrs}")

    # 4. Bar count estimate
    symbol_list = ", ".join(f"'{s}'" for s in crypto_symbols)
    bar_count = client.command(
        f"SELECT count() FROM opendeviationbar_cache.open_deviation_bars WHERE symbol IN ({symbol_list})"
    )
    _info(f"Total crypto bars to be dropped: {bar_count:,}")

    # 5. Check for running services (informational)
    _header("Service Status Check")
    _warn("Cannot check systemd services remotely. On bigblack, verify:")
    print("    systemctl --user is-active opendeviationbar-sidecar")
    print("    systemctl --user is-active opendeviationbar-kintsugi")
    print("    systemctl --user is-active opendeviationbar-kintsugi-catchup")

    # 6. Instructions
    _header("Next Steps")
    print(f"""
  To execute the migration:

  1. SSH to bigblack and stop services:
     {YELLOW}systemctl --user stop opendeviationbar-sidecar opendeviationbar-kintsugi opendeviationbar-kintsugi-catchup{RESET}

  2. Run the execute phase:
     {YELLOW}python scripts/fleet_recompute_aion.py --execute{RESET}

  3. After execute, start kintsugi catch-up:
     {YELLOW}systemctl --user start opendeviationbar-kintsugi-catchup{RESET}

  4. Monitor and verify:
     {YELLOW}python scripts/fleet_recompute_aion.py --verify{RESET}
""")

    return 0


# =============================================================================
# Phase 2: Execute
# =============================================================================

def run_execute(
    client: clickhouse_connect.driver.Client,
    *,
    dry_run: bool = False,
) -> int:
    """Execute the Aion migration: DROP PARTITION + wipe telemetry."""
    _header("Phase 2: Execute Migration")

    if dry_run:
        _warn("DRY RUN -- no changes will be made")

    # 1. Get crypto symbols and partitions
    crypto_symbols = _get_crypto_symbols()
    partitions = _discover_crypto_partitions(client, crypto_symbols)

    if not partitions:
        _warn("No crypto partitions found -- nothing to do")
        return 0

    _info(f"Found {len(partitions)} crypto partitions to drop")

    # 2. Confirmation prompt (unless dry-run)
    if not dry_run:
        symbol_list = ", ".join(f"'{s}'" for s in crypto_symbols)
        bar_count = client.command(
            f"SELECT count() FROM opendeviationbar_cache.open_deviation_bars WHERE symbol IN ({symbol_list})"
        )
        print(f"\n  {RED}{BOLD}WARNING: This will permanently delete {bar_count:,} bars across {len(partitions)} partitions.{RESET}")
        print(f"  {RED}This action cannot be undone. Kintsugi must recompute all bars from scratch.{RESET}")
        print()
        confirm = input(f"  Type {BOLD}YES{RESET} to proceed: ")
        if confirm != "YES":
            _err("Aborted by user")
            return 1

    # 3. DROP PARTITION for each crypto partition
    _header("Step 1: DROP PARTITION")
    dropped = 0
    failed = 0

    for symbol, threshold in partitions:
        stmt = (
            f"ALTER TABLE opendeviationbar_cache.open_deviation_bars "
            f"DROP PARTITION ('{symbol}', {threshold})"
        )
        if dry_run:
            _info(f"[DRY RUN] Would execute: {stmt}")
            dropped += 1
        else:
            try:
                client.command(stmt)
                _ok(f"Dropped partition ({symbol}, {threshold})")
                dropped += 1
            except Exception as e:
                _err(f"Failed to drop ({symbol}, {threshold}): {e}")
                failed += 1

    _info(f"Dropped: {dropped}, Failed: {failed}")

    # 4. Verify partitions are gone
    if not dry_run:
        remaining = _discover_crypto_partitions(client, crypto_symbols)
        if remaining:
            _warn(f"{len(remaining)} crypto partitions still exist after DROP")
            for sym, thr in remaining:
                print(f"    ({sym}, {thr})")
        else:
            _ok("All crypto partitions confirmed dropped")

    # 5. Wipe telemetry tables
    _header("Step 2: Wipe Telemetry")
    symbol_sql = ", ".join(f"'{s}'" for s in crypto_symbols)

    telemetry_tables = [
        "opendeviationbar_cache.kintsugi_log",
        "opendeviationbar_cache.population_checkpoints",
        "opendeviationbar_cache.backfill_requests",
    ]

    for table in telemetry_tables:
        stmt = (
            f"ALTER TABLE {table} DELETE WHERE symbol IN ({symbol_sql}) "
            f"SETTINGS mutations_sync = 1"
        )
        if dry_run:
            _info(f"[DRY RUN] Would execute: ALTER TABLE {table} DELETE WHERE symbol IN (...)")
        else:
            try:
                client.command(stmt)
                _ok(f"Wiped crypto entries from {table}")
            except Exception as e:
                # Table may not exist on all setups
                _warn(f"Could not wipe {table}: {e}")

    # 6. Wipe local heartbeat state + checkpoint files
    _header("Step 3: Wipe Local State Files")
    state_files = [
        pathlib.Path.home() / ".cache" / "opendeviationbar" / "heartbeat" / "stathera_counts.json",
    ]

    for fpath in state_files:
        if dry_run:
            _info(f"[DRY RUN] Would delete: {fpath}")
        elif fpath.exists():
            fpath.unlink()
            _ok(f"Deleted {fpath}")
        else:
            _info(f"Already absent: {fpath}")

    # Wipe crypto checkpoint files
    checkpoint_dir = pathlib.Path.home() / ".cache" / "opendeviationbar" / "checkpoints"
    if checkpoint_dir.exists():
        crypto_set = set(crypto_symbols)
        for cp_file in checkpoint_dir.iterdir():
            if cp_file.is_file():
                # Checkpoint files are named like BTCUSDT_250.json
                stem = cp_file.stem
                # Check if any crypto symbol is in the filename
                if any(sym in stem for sym in crypto_set):
                    if dry_run:
                        _info(f"[DRY RUN] Would delete: {cp_file}")
                    else:
                        cp_file.unlink()
                        _ok(f"Deleted checkpoint: {cp_file.name}")

    # 7. Summary
    _header("Execution Complete")
    ts = datetime.now(tz=UTC).strftime("%Y-%m-%d %H:%M:%S UTC")
    print(f"\n  Timestamp: {ts}")
    print(f"  Partitions dropped: {dropped}")
    print(f"  Partitions failed: {failed}")

    if not dry_run:
        print(f"""
  {BOLD}Next steps:{RESET}

  1. Start kintsugi catch-up to recompute all Aion bars:
     {YELLOW}systemctl --user start opendeviationbar-kintsugi-catchup{RESET}

  2. Monitor progress:
     {YELLOW}journalctl --user -u opendeviationbar-kintsugi-catchup -f{RESET}

  3. After kintsugi completes (~1-3h), verify:
     {YELLOW}python scripts/fleet_recompute_aion.py --verify{RESET}

  4. Restart sidecar:
     {YELLOW}systemctl --user start opendeviationbar-sidecar{RESET}
""")

    return 0 if failed == 0 else 1


# =============================================================================
# Phase 3: Verify
# =============================================================================

def run_verify(client: clickhouse_connect.driver.Client) -> int:
    """Post-recompute Stathera verification audit."""
    _header("Phase 3: Stathera Verification Audit")

    crypto_symbols = _get_crypto_symbols()
    symbol_sql = ", ".join(f"'{s}'" for s in crypto_symbols)

    # Check if any bars exist yet
    bar_count = client.command(
        f"SELECT count() FROM opendeviationbar_cache.open_deviation_bars "
        f"WHERE symbol IN ({symbol_sql})"
    )
    _info(f"Total crypto bars in ClickHouse: {bar_count:,}")

    if bar_count == 0:
        _warn("No crypto bars found -- kintsugi recompute may not have started yet")
        return 1

    # Run Stathera audit: detect gaps and overlaps via trade-ID continuity
    _header("Stathera Audit (Trade-ID Continuity)")

    audit_query = f"""
    SELECT symbol, threshold_decimal_bps,
           countIf(tid_delta < 0) AS overlaps,
           countIf(tid_delta > 0) AS gaps
    FROM (
        SELECT symbol, threshold_decimal_bps,
               first_agg_trade_id - lagInFrame(last_agg_trade_id, 1)
                   OVER (PARTITION BY symbol, threshold_decimal_bps
                         ORDER BY first_agg_trade_id, close_time_us) - 1 AS tid_delta,
               row_number() OVER (PARTITION BY symbol, threshold_decimal_bps
                                  ORDER BY first_agg_trade_id, close_time_us) AS rn
        FROM opendeviationbar_cache.open_deviation_bars FINAL
        WHERE symbol IN ({symbol_sql})
          AND first_agg_trade_id > 0 AND last_agg_trade_id > 0
    )
    WHERE rn > 1 AND tid_delta != 0
    GROUP BY symbol, threshold_decimal_bps
    ORDER BY symbol, threshold_decimal_bps
    """

    result = client.query(audit_query)
    total_gaps = 0
    total_overlaps = 0

    if result.result_rows:
        print(f"\n  {'Symbol':<16} {'Threshold':>10} {'Gaps':>8} {'Overlaps':>10}")
        print(f"  {'------':<16} {'---------':>10} {'----':>8} {'--------':>10}")
        for row in result.result_rows:
            symbol, threshold, overlaps, gaps = row
            total_gaps += gaps
            total_overlaps += overlaps
            status = RED if (gaps > 0 or overlaps > 0) else GREEN
            print(f"  {status}{symbol:<16} {threshold:>10} {gaps:>8} {overlaps:>10}{RESET}")
    else:
        _ok("No Stathera anomalies detected (0 gaps, 0 overlaps)")

    # Per-symbol bar count summary
    _header("Bar Count Summary")
    count_query = f"""
    SELECT symbol, threshold_decimal_bps, count() AS bars,
           min(close_time_us) AS earliest_us,
           max(close_time_us) AS latest_us
    FROM opendeviationbar_cache.open_deviation_bars FINAL
    WHERE symbol IN ({symbol_sql})
    GROUP BY symbol, threshold_decimal_bps
    ORDER BY symbol, threshold_decimal_bps
    """
    count_result = client.query(count_query)

    print(f"\n  {'Symbol':<16} {'Threshold':>10} {'Bars':>12} {'Earliest':>12} {'Latest':>12}")
    print(f"  {'------':<16} {'---------':>10} {'----':>12} {'--------':>12} {'------':>12}")
    for row in count_result.result_rows:
        symbol, threshold, bars, earliest_us, latest_us = row
        earliest_dt = datetime.fromtimestamp(earliest_us / 1_000_000, tz=UTC).strftime("%Y-%m-%d") if earliest_us else "N/A"
        latest_dt = datetime.fromtimestamp(latest_us / 1_000_000, tz=UTC).strftime("%Y-%m-%d") if latest_us else "N/A"
        print(f"  {symbol:<16} {threshold:>10} {bars:>12,} {earliest_dt:>12} {latest_dt:>12}")

    # Check ouroboros_mode distribution
    _header("Ouroboros Mode Distribution")
    mode_query = f"""
    SELECT ouroboros_mode, count() AS bars
    FROM opendeviationbar_cache.open_deviation_bars FINAL
    WHERE symbol IN ({symbol_sql})
    GROUP BY ouroboros_mode
    ORDER BY bars DESC
    """
    mode_result = client.query(mode_query)
    for row in mode_result.result_rows:
        mode, bars = row
        status = GREEN if mode == "aion" else YELLOW
        print(f"  {status}{mode}: {bars:,} bars{RESET}")

    # Final verdict
    _header("Verdict")
    if total_gaps == 0 and total_overlaps == 0:
        _ok(f"CLEAN: 0 gaps, 0 overlaps across all crypto (symbol, threshold) pairs")
        _ok(f"Total crypto bars: {bar_count:,}")
        return 0
    else:
        _err(f"ISSUES FOUND: {total_gaps} gaps, {total_overlaps} overlaps")
        _err("Kintsugi may still be running -- wait and re-verify")
        return 1


# =============================================================================
# CLI
# =============================================================================

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fleet Recompute: Aion Migration for ClickHouse (one-time v7.0 tool)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python scripts/fleet_recompute_aion.py --preflight   # Discover partitions, show plan
  python scripts/fleet_recompute_aion.py --execute      # Drop partitions + wipe telemetry
  python scripts/fleet_recompute_aion.py --execute --dry-run  # Show what would happen
  python scripts/fleet_recompute_aion.py --verify       # Post-recompute Stathera audit
        """,
    )

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "--preflight",
        action="store_true",
        help="Phase 1: Discover partitions and show migration plan",
    )
    group.add_argument(
        "--execute",
        action="store_true",
        help="Phase 2: DROP PARTITION and wipe telemetry (destructive!)",
    )
    group.add_argument(
        "--verify",
        action="store_true",
        help="Phase 3: Post-recompute Stathera audit",
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be done without making changes (--execute only)",
    )

    return parser.parse_args()


def main() -> int:
    args = parse_args()

    # Connect to ClickHouse
    _header("Connecting to ClickHouse")
    try:
        client, tunnel = connect()
        _ok("Connected to ClickHouse")
    except Exception as e:
        _err(f"Failed to connect: {e}")
        return 2

    try:
        if args.preflight:
            return run_preflight(client)
        elif args.execute:
            return run_execute(client, dry_run=args.dry_run)
        elif args.verify:
            return run_verify(client)
        else:
            _err("No phase specified")
            return 1
    finally:
        if tunnel is not None:
            tunnel.stop()


if __name__ == "__main__":
    sys.exit(main())
