#!/usr/bin/env python3
# ruff: noqa: DTZ001
"""Memray baseline memory profile for the backfill pipeline.

Profiles the backfill pipeline using local fixture data (no network required).
Two modes:
  - batch: loads all trades as Python dicts, then processes (old code path)
  - streaming: Polars → Arrow → process_trades_arrow() (production path)

Usage:
    # Run with memray native tracking (sees Python + Rust allocations)
    .venv/bin/python3 -m memray run --native -o /tmp/odb_profile.bin \
        scripts/profile_memory_baseline.py

    # Generate flamegraph
    .venv/bin/python3 -m memray flamegraph /tmp/odb_profile.bin -o /tmp/odb_flamegraph.html

    # Generate summary (top allocators)
    .venv/bin/python3 -m memray summary /tmp/odb_profile.bin

    # Run without memray (uses tracemalloc + psutil for quick baseline)
    python scripts/profile_memory_baseline.py

Based on: scripts/validate_memory_efficiency.py (Issue #248)
"""

from __future__ import annotations

import argparse
import gc
import os
import resource
import sys
import time
import tracemalloc
from pathlib import Path

FIXTURE_DIR = Path(__file__).parent.parent / "tests" / "fixtures"

# Fixtures available for profiling (symbol, zip_filename, approx_trades)
FIXTURES = [
    ("BTCUSDT", "BTCUSDT-aggTrades-2024-01-01.zip", 863_000),
    ("ETHUSDT", "ETHUSDT-aggTrades-2024-01-01.zip", 400_000),
    ("SOLUSDT", "SOLUSDT-aggTrades-2024-01-01.zip", 300_000),
]


def get_rss_gb() -> float:
    """Get current RSS (resident set size) in GB."""
    try:
        import psutil
        return psutil.Process().memory_info().rss / (1024**3)
    except ImportError:
        # Fallback: ru_maxrss (peak, not current)
        usage = resource.getrusage(resource.RUSAGE_SELF)
        if sys.platform == "darwin":
            return usage.ru_maxrss / (1024**3)  # macOS: bytes
        return usage.ru_maxrss / (1024**2)  # Linux: KB


def _parse_csv_rows(zip_path: Path, limit: int | None = None) -> list[tuple]:
    """Parse CSV rows from ZIP as tuples (shared by both modes)."""
    import zipfile

    rows = []
    with zipfile.ZipFile(zip_path) as zf:
        csv_name = zf.namelist()[0]
        with zf.open(csv_name) as f:
            for line in f:
                parts = line.decode().strip().split(",")
                if len(parts) >= 7:
                    rows.append((
                        int(parts[0]),    # agg_trade_id
                        float(parts[1]),  # price
                        float(parts[2]),  # quantity
                        int(parts[3]),    # first_trade_id
                        int(parts[4]),    # last_trade_id
                        int(parts[5]),    # timestamp (ms)
                        parts[6].strip().lower() == "true",  # is_buyer_maker
                    ))
                    if limit and len(rows) >= limit:
                        break
    return rows


def profile_process_trades(symbol: str, zip_path: Path, threshold_dbps: int = 250) -> dict:
    """Profile process_trades() with trades as Python dicts (old code path)."""
    from opendeviationbar import OpenDeviationBarProcessor

    gc.collect()
    rss_before = get_rss_gb()

    print(f"  Loading trades from {zip_path.name}...")
    raw_rows = _parse_csv_rows(zip_path)
    trades = [
        {
            "agg_trade_id": r[0], "price": r[1], "quantity": r[2],
            "first_trade_id": r[3], "last_trade_id": r[4],
            "timestamp": r[5], "is_buyer_maker": r[6],
        }
        for r in raw_rows
    ]
    del raw_rows
    trade_count = len(trades)

    rss_after_load = get_rss_gb()
    print(f"  Loaded {trade_count:,} trades (RSS: {rss_before:.2f} -> {rss_after_load:.2f} GB)")

    print(f"  Processing via process_trades() (threshold={threshold_dbps} dbps)...")
    processor = OpenDeviationBarProcessor(threshold_decimal_bps=threshold_dbps)

    start = time.time()
    chunk_size = 100_000
    total_bars = 0
    for i in range(0, len(trades), chunk_size):
        chunk = trades[i : i + chunk_size]
        bars = processor.process_trades(chunk)
        total_bars += len(bars)

    elapsed = time.time() - start
    rss_after_process = get_rss_gb()

    print(f"  Generated {total_bars} bars in {elapsed:.2f}s")
    print(f"  RSS after processing: {rss_after_process:.2f} GB")

    del trades
    gc.collect()
    rss_after_cleanup = get_rss_gb()

    return {
        "symbol": symbol,
        "threshold": threshold_dbps,
        "total_trades": trade_count,
        "total_bars": total_bars,
        "elapsed_s": elapsed,
        "rss_before_gb": rss_before,
        "rss_after_load_gb": rss_after_load,
        "rss_after_process_gb": rss_after_process,
        "rss_after_cleanup_gb": rss_after_cleanup,
        "rss_peak_gb": max(rss_after_load, rss_after_process),
    }


def profile_populate_pipeline(symbol: str, zip_path: Path) -> dict:
    """Profile the actual production Arrow pipeline.

    Simulates: CSV → tuples → Polars → Arrow → process_trades_arrow() → bars.
    This is the real code path used by populate_cache_resumable().
    """
    import polars as pl
    from opendeviationbar import OpenDeviationBarProcessor

    gc.collect()
    rss_before = get_rss_gb()
    start = time.time()

    processor = OpenDeviationBarProcessor(threshold_decimal_bps=250)
    total_bars = 0
    total_trades = 0
    chunk_size = 100_000

    print("  Streaming CSV → Polars → Arrow → process_trades_arrow()...")
    raw_rows = _parse_csv_rows(zip_path)
    rss_after_parse = get_rss_gb()
    print(f"  Parsed {len(raw_rows):,} rows (RSS: {rss_before:.3f} -> {rss_after_parse:.3f} GB)")

    # Process in chunks through Arrow path (production pipeline)
    for i in range(0, len(raw_rows), chunk_size):
        chunk = raw_rows[i : i + chunk_size]
        df = pl.DataFrame({
            "agg_trade_id": [r[0] for r in chunk],
            "price": [r[1] for r in chunk],
            "quantity": [r[2] for r in chunk],
            "first_trade_id": [r[3] for r in chunk],
            "last_trade_id": [r[4] for r in chunk],
            "timestamp": [r[5] for r in chunk],
            "is_buyer_maker": [r[6] for r in chunk],
        })
        arrow_batch = df.to_arrow()
        bars_arrow = processor.process_trades_arrow(arrow_batch)
        bars_df = pl.from_arrow(bars_arrow)
        total_bars += len(bars_df)
        total_trades += len(chunk)
        del df, arrow_batch, bars_arrow, bars_df

    del raw_rows
    elapsed = time.time() - start
    rss_after = get_rss_gb()

    gc.collect()
    rss_after_cleanup = get_rss_gb()

    return {
        "symbol": symbol,
        "total_trades": total_trades,
        "total_bars": total_bars,
        "elapsed_s": elapsed,
        "rss_before_gb": rss_before,
        "rss_peak_gb": rss_after,
        "rss_after_cleanup_gb": rss_after_cleanup,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Memray baseline memory profiler")
    parser.add_argument(
        "--mode",
        choices=["batch", "streaming", "both"],
        default="both",
        help="Profile mode: batch (Python dicts), streaming (Arrow path), both",
    )
    parser.add_argument(
        "--symbol",
        default="BTCUSDT",
        help="Symbol to profile (default: BTCUSDT)",
    )
    args = parser.parse_args()

    print("=" * 70)
    print("opendeviationbar Memory Profile Baseline")
    print("=" * 70)

    try:
        from opendeviationbar import __version__
        print(f"Version: {__version__}")
    except ImportError as e:
        print(f"ERROR: {e}")
        print("Run: maturin develop")
        return 1

    # Find fixture
    zip_path = FIXTURE_DIR / f"{args.symbol}-aggTrades-2024-01-01.zip"
    if not zip_path.exists():
        print(f"ERROR: Fixture not found: {zip_path}")
        print(f"Available: {[f.name for f in FIXTURE_DIR.glob('*.zip')]}")
        return 1

    print(f"Fixture: {zip_path.name} ({zip_path.stat().st_size / 1024 / 1024:.1f} MB)")
    print(f"Initial RSS: {get_rss_gb():.3f} GB")

    # Detect if running under memray
    under_memray = "MEMRAY" in os.environ or any("memray" in arg for arg in sys.argv)
    if not under_memray:
        tracemalloc.start()
        print("(Not running under memray — using tracemalloc fallback)")
    print()

    results = []

    if args.mode in ("batch", "both"):
        print("-" * 70)
        print("Mode 1: BATCH (Python dicts → process_trades())")
        print("  Simulates: old code path with Python dict intermediaries")
        print("-" * 70)
        result = profile_process_trades(args.symbol, zip_path)
        results.append(("batch", result))
        print()

    if args.mode in ("streaming", "both"):
        print("-" * 70)
        print("Mode 2: ARROW (Polars → Arrow → process_trades_arrow())")
        print("  Simulates: v13.1+ production pipeline (zero-copy)")
        print("-" * 70)
        result = profile_populate_pipeline(args.symbol, zip_path)
        results.append(("arrow", result))
        print()

    # Summary
    print("=" * 70)
    print("RESULTS SUMMARY")
    print("=" * 70)
    for mode, r in results:
        print(f"\n  {mode.upper()}:")
        print(f"    Trades:     {r.get('total_trades', 0):>12,}")
        print(f"    Bars:       {r.get('total_bars', 0):>12,}")
        print(f"    Time:       {r['elapsed_s']:>12.2f} s")
        print(f"    RSS before: {r['rss_before_gb']:>12.3f} GB")
        peak = r.get("rss_peak_gb", r.get("rss_after_process_gb", 0))
        print(f"    RSS peak:   {peak:>12.3f} GB")
        print(f"    RSS after:  {r.get('rss_after_cleanup_gb', 0):>12.3f} GB")

    if not under_memray:
        current, peak = tracemalloc.get_traced_memory()
        print(f"\n  tracemalloc peak: {peak / 1024 / 1024:.1f} MB")
        tracemalloc.stop()

    peak_rss = max(
        r.get("rss_peak_gb", r.get("rss_after_process_gb", 0))
        for _, r in results
    )
    print(f"\n  Overall peak RSS: {peak_rss:.3f} GB")
    print()

    if under_memray:
        print("Next steps:")
        print("  .venv/bin/python3 -m memray flamegraph /tmp/odb_profile.bin -o /tmp/odb_flamegraph.html")
        print("  .venv/bin/python3 -m memray summary /tmp/odb_profile.bin")

    return 0


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