#!/usr/bin/env python3
"""Reseed Parquet tick cache files from millisecond to microsecond timestamps.

One-off migration script that walks all Parquet tick files in the local cache
and upscales ms timestamps to us (multiply by 1000). Files that already have
microsecond timestamps (>= 1e13) are skipped.

This is a direct file-level operation: read -> transform -> atomic write.
Does NOT use TickStorage.write_ticks() to avoid dedup/partitioning logic.

Usage:
    # Dry run (report only)
    python scripts/reseed_parquet_to_microseconds.py --dry-run

    # Reseed all symbols
    python scripts/reseed_parquet_to_microseconds.py

    # Reseed single symbol
    python scripts/reseed_parquet_to_microseconds.py --symbol BTCUSDT

    # Verify reseeded data produces valid bars
    python scripts/reseed_parquet_to_microseconds.py --verify --symbol BTCUSDT
"""
from __future__ import annotations

import argparse
import logging
import os
import sys
import tempfile
from pathlib import Path

import polars as pl

from opendeviationbar.storage.parquet_io import (
    COMPRESSION,
    COMPRESSION_LEVEL,
    get_cache_dir,
)

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

# Microsecond/millisecond detection threshold (from parquet.py).
# ms values for modern dates are ~1e12; us values are ~1e15.
# Threshold 1e13 safely separates the two for any date 2001+.
_US_MS_THRESHOLD = 10_000_000_000_000  # 1e13


def _atomic_write_parquet(df: pl.DataFrame, target_path: Path) -> None:
    """Write Parquet atomically: write to temp, fsync, rename."""
    fd, temp_path_str = tempfile.mkstemp(
        dir=target_path.parent,
        prefix=".reseed_",
        suffix=".tmp",
    )
    temp_path = Path(temp_path_str)
    try:
        os.close(fd)
        df.write_parquet(
            temp_path,
            compression=COMPRESSION,
            compression_level=COMPRESSION_LEVEL,
        )
        # Sync to disk before rename
        with temp_path.open("rb") as f:
            os.fsync(f.fileno())
        # Atomic rename (POSIX)
        os.replace(str(temp_path), str(target_path))
    except (OSError, RuntimeError):
        if temp_path.exists():
            temp_path.unlink(missing_ok=True)
        raise


def reseed_file(file_path: Path, *, dry_run: bool = False) -> str:
    """Reseed a single Parquet file from ms to us timestamps.

    Returns
    -------
    str
        "reseeded", "already_us", or "skipped" (empty/no timestamp col).
    """
    try:
        df = pl.read_parquet(file_path)
    except Exception:
        logger.warning("Failed to read %s, skipping", file_path)
        return "skipped"

    if df.is_empty():
        return "skipped"

    if "timestamp" not in df.columns:
        logger.warning("No timestamp column in %s, skipping", file_path)
        return "skipped"

    max_val = df["timestamp"].max()
    if max_val is None:
        return "skipped"

    if max_val >= _US_MS_THRESHOLD:
        # Already microsecond timestamps
        return "already_us"

    # Needs reseeding: multiply by 1000
    if dry_run:
        return "reseeded"

    original_rows = len(df)
    df = df.with_columns((pl.col("timestamp") * 1000).alias("timestamp"))

    # Corruption guard: verify row count matches
    if len(df) != original_rows:
        logger.error(
            "Row count mismatch after transform: %d vs %d for %s",
            len(df), original_rows, file_path,
        )
        return "skipped"

    _atomic_write_parquet(df, file_path)
    return "reseeded"


def reseed_all(
    *,
    dry_run: bool = False,
    symbol_filter: str | None = None,
) -> dict[str, dict[str, int]]:
    """Walk all Parquet tick files and reseed ms -> us.

    Returns per-symbol summary: {symbol: {reseeded, already_us, skipped, total}}.
    """
    ticks_dir = get_cache_dir() / "ticks"
    if not ticks_dir.exists():
        logger.warning("Ticks directory does not exist: %s", ticks_dir)
        return {}

    summaries: dict[str, dict[str, int]] = {}

    # Discover symbol directories
    symbol_dirs = sorted(
        d for d in ticks_dir.iterdir()
        if d.is_dir() and d.name != ""
    )

    if symbol_filter:
        symbol_dirs = [d for d in symbol_dirs if d.name == symbol_filter]
        if not symbol_dirs:
            logger.warning("No directory found for symbol %s", symbol_filter)
            return {}

    for symbol_dir in symbol_dirs:
        symbol = symbol_dir.name
        files = sorted(symbol_dir.glob("*.parquet"))

        # Skip .tmp files (leftover from interrupted writes)
        files = [f for f in files if not f.name.endswith(".tmp")]

        counts = {"reseeded": 0, "already_us": 0, "skipped": 0, "total": len(files)}

        for file_path in files:
            result = reseed_file(file_path, dry_run=dry_run)
            counts[result] += 1

            if result == "reseeded":
                n_rows = 0
                try:
                    if not dry_run:
                        n_rows = len(pl.read_parquet(file_path))
                except Exception:
                    pass
                action = "would reseed" if dry_run else "reseeded ms->us"
                logger.info(
                    "%s %s: %s (%d rows)",
                    symbol, file_path.stem, action, n_rows,
                )

        summaries[symbol] = counts

    # Print summary table
    print("\n" + "=" * 70)
    print(f"{'Symbol':<30} {'Reseeded':>10} {'Already us':>12} {'Skipped':>10} {'Total':>8}")
    print("-" * 70)
    for symbol, counts in sorted(summaries.items()):
        print(
            f"{symbol:<30} {counts['reseeded']:>10} {counts['already_us']:>12} "
            f"{counts['skipped']:>10} {counts['total']:>8}"
        )
    print("=" * 70)

    mode = "DRY RUN" if dry_run else "EXECUTED"
    total_reseeded = sum(c["reseeded"] for c in summaries.values())
    total_files = sum(c["total"] for c in summaries.values())
    print(f"\n{mode}: {total_reseeded}/{total_files} files reseeded across {len(summaries)} symbols\n")

    return summaries


def verify_reseeded_data(symbol: str = "BTCUSDT") -> bool:
    """Verify reseeded data produces valid bars with Stathera continuity.

    Picks test days (pre-2025, boundary, post-2025), verifies:
    1. All timestamps are microseconds (>= _US_MS_THRESHOLD)
    2. Timestamp column is Int64
    3. Bars can be produced via Rust processor (us-native Arrow path)
    4. Intra-day Stathera continuity holds
    5. Cross-day Stathera continuity at 2024-12-31 / 2025-01-01 boundary

    Uses process_trades_arrow_native() which accepts microsecond timestamps
    directly, validating the full us-native pipeline end-to-end.

    Returns True if all checks pass, False otherwise.
    """
    from opendeviationbar._core import PyOpenDeviationBarProcessor

    ticks_dir = get_cache_dir() / "ticks"

    # Find the symbol directory (try both bare and prefixed names)
    symbol_dir = ticks_dir / symbol
    if not symbol_dir.exists():
        symbol_dir = ticks_dir / f"BINANCE_SPOT_{symbol}"
    if not symbol_dir.exists():
        logger.warning(
            "No Parquet directory for %s at %s, skipping verification", symbol, ticks_dir
        )
        return True  # Not a failure — just no data to verify

    # Pick test days
    test_days = ["2024-06-15", "2024-12-31", "2025-01-01"]
    results: dict[str, dict] = {}
    all_pass = True

    for day in test_days:
        parquet_path = symbol_dir / f"{day}.parquet"
        if not parquet_path.exists():
            logger.warning("  %s: file not found, skipping", day)
            results[day] = {"status": "not_found"}
            continue

        df = pl.read_parquet(parquet_path)
        n_rows = len(df)

        if "timestamp" not in df.columns:
            logger.error("  %s: no timestamp column", day)
            results[day] = {"status": "FAIL", "reason": "no timestamp column"}
            all_pass = False
            continue

        # Check 1: All timestamps >= _US_MS_THRESHOLD
        min_ts = df["timestamp"].min()
        if min_ts is not None and min_ts < _US_MS_THRESHOLD:
            logger.error(
                "  %s: has ms timestamps (min=%d < %d)",
                day, min_ts, _US_MS_THRESHOLD,
            )
            results[day] = {"status": "FAIL", "reason": f"ms timestamps (min={min_ts})"}
            all_pass = False
            continue

        # Check 2: dtype is Int64
        ts_dtype = df.schema["timestamp"]
        if ts_dtype != pl.Int64:
            logger.error("  %s: timestamp dtype is %s, expected Int64", day, ts_dtype)
            results[day] = {"status": "FAIL", "reason": f"wrong dtype {ts_dtype}"}
            all_pass = False
            continue

        # Check 3: Process through Rust via us-native Arrow path
        # process_trades_arrow_native() accepts microsecond timestamps directly,
        # validating the full us pipeline without any ms conversion.
        import pyarrow as pa

        # Select columns needed by Rust
        tick_cols = ["agg_trade_id", "price", "quantity", "timestamp", "is_buyer_maker"]
        select_cols = [c for c in tick_cols if c in df.columns]
        df_for_rust = df.select(select_cols)

        trades_arrow = df_for_rust.to_arrow()
        batches = trades_arrow.to_batches()
        if not batches:
            logger.warning("  %s: no Arrow batches", day)
            results[day] = {"status": "FAIL", "reason": "empty Arrow batch"}
            all_pass = False
            continue

        if len(batches) == 1:
            batch = batches[0]
        else:
            batch = pa.Table.from_batches(batches).combine_chunks().to_batches()[0]

        processor = PyOpenDeviationBarProcessor(250)
        # Use us-native path: timestamps are already microseconds
        try:
            bars_arrow = processor.process_trades_arrow_native(batch)
        except AttributeError:
            # Fallback: if process_trades_arrow_native not available in this build,
            # convert to ms and use the standard path
            logger.warning("  process_trades_arrow_native not available, falling back to ms path")
            df_ms = df.with_columns((pl.col("timestamp") // 1000).alias("timestamp"))
            df_ms = df_ms.select(select_cols)
            trades_arrow_ms = df_ms.to_arrow()
            batches_ms = trades_arrow_ms.to_batches()
            batch_ms = batches_ms[0] if len(batches_ms) == 1 else pa.Table.from_batches(batches_ms).combine_chunks().to_batches()[0]
            bars_arrow = processor.process_trades_arrow(batch_ms)
        bars_df = pl.from_arrow(bars_arrow)

        # Get orphan bar
        orphan = processor.reset_at_ouroboros()

        n_bars = len(bars_df)
        if n_bars == 0:
            logger.warning("  %s: 0 bars produced (may be normal for low-volume days)", day)
            results[day] = {"status": "OK", "n_rows": n_rows, "n_bars": 0, "stathera": "n/a"}
            continue

        # Check 4: Intra-day Stathera continuity
        stathera_ok = True
        if n_bars >= 2 and "last_agg_trade_id" in bars_df.columns and "first_agg_trade_id" in bars_df.columns:
            last_ids = bars_df["last_agg_trade_id"][:-1].to_list()
            first_ids = bars_df["first_agg_trade_id"][1:].to_list()
            for i, (last_id, first_id) in enumerate(zip(last_ids, first_ids)):
                if last_id + 1 != first_id:
                    logger.error(
                        "  %s: Stathera break at bar %d: last_agg_trade_id=%d, next first=%d",
                        day, i, last_id, first_id,
                    )
                    stathera_ok = False
                    all_pass = False
                    break

        status = "OK" if stathera_ok else "FAIL (Stathera)"
        results[day] = {
            "status": status,
            "n_rows": n_rows,
            "n_bars": n_bars,
            "stathera": "OK" if stathera_ok else "FAIL",
        }

    # Check 5: Cross-day continuity (2024-12-31 -> 2025-01-01)
    cross_day_ok = "n/a"
    day1_path = symbol_dir / "2024-12-31.parquet"
    day2_path = symbol_dir / "2025-01-01.parquet"

    if day1_path.exists() and day2_path.exists():
        tick_cols = ["agg_trade_id", "price", "quantity", "timestamp", "is_buyer_maker"]

        def _process_day_for_verify(path: Path) -> tuple:
            """Read Parquet and process through Rust, returning (bars_df, processor)."""
            df = pl.read_parquet(path)
            select_cols = [c for c in tick_cols if c in df.columns]
            df_sel = df.select(select_cols)
            batches = df_sel.to_arrow().to_batches()
            if len(batches) == 1:
                batch = batches[0]
            else:
                import pyarrow as pa
                batch = pa.Table.from_batches(batches).combine_chunks().to_batches()[0]
            proc = PyOpenDeviationBarProcessor(250)
            try:
                bars = pl.from_arrow(proc.process_trades_arrow_native(batch))
            except AttributeError:
                # Fallback to ms path
                df_ms = df.with_columns((pl.col("timestamp") // 1000).alias("timestamp"))
                df_ms = df_ms.select(select_cols)
                b_ms = df_ms.to_arrow().to_batches()
                batch_ms = b_ms[0] if len(b_ms) == 1 else pa.Table.from_batches(b_ms).combine_chunks().to_batches()[0]
                bars = pl.from_arrow(proc.process_trades_arrow(batch_ms))
            return bars, proc

        bars1, p1 = _process_day_for_verify(day1_path)
        orphan1 = p1.reset_at_ouroboros()

        bars2, p2 = _process_day_for_verify(day2_path)

        # Check Stathera at boundary
        if not bars1.is_empty() and not bars2.is_empty():
            # Include orphan bar if it exists
            if orphan1 is not None:
                orphan_df = pl.from_arrow(orphan1) if hasattr(orphan1, '__len__') else None
                if orphan_df is not None and not orphan_df.is_empty():
                    bars1 = pl.concat([bars1, orphan_df])

            if "last_agg_trade_id" in bars1.columns and "first_agg_trade_id" in bars2.columns:
                last_of_day1 = bars1["last_agg_trade_id"][-1]
                first_of_day2 = bars2["first_agg_trade_id"][0]
                if last_of_day1 + 1 == first_of_day2:
                    cross_day_ok = "OK"
                else:
                    cross_day_ok = f"FAIL (gap: {last_of_day1} -> {first_of_day2})"
                    all_pass = False
        else:
            cross_day_ok = "skipped (empty bars)"

    # Print verification report
    print("\nVerification Report:")
    for day, info in results.items():
        if info["status"] == "not_found":
            print(f"  {day}: file not found (skipped)")
        elif "n_rows" in info:
            print(
                f"  {day}: {info['n_rows']} ticks, all us, "
                f"{info['n_bars']} bars, Stathera {info['stathera']}"
            )
        else:
            print(f"  {day}: {info['status']} - {info.get('reason', '')}")

    print(f"  Cross-day (2024-12-31 -> 2025-01-01): {cross_day_ok}")
    print()

    if all_pass:
        print("VERIFICATION PASSED")
    else:
        print("VERIFICATION FAILED")

    return all_pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Reseed Parquet tick files from ms to us timestamps"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Report only, no modifications",
    )
    parser.add_argument(
        "--symbol",
        type=str,
        help="Restrict to single symbol",
    )
    parser.add_argument(
        "--verify",
        action="store_true",
        help="Verify reseeded data produces valid bars with Stathera continuity",
    )
    args = parser.parse_args()

    if args.verify:
        symbol = args.symbol or "BTCUSDT"
        ok = verify_reseeded_data(symbol)
        sys.exit(0 if ok else 1)
    else:
        reseed_all(dry_run=args.dry_run, symbol_filter=args.symbol)
