#!/usr/bin/env python3
"""
Memory Usage Test Script

Generates a large wasot file and verifies that processing it uses bounded memory.
Uses DHAT profiler to measure peak heap allocation.

Usage:
    python scripts/test_memory_usage.py
    python scripts/test_memory_usage.py --size 200MB --max-memory 10MB
    python scripts/test_memory_usage.py --skip-generate  # reuse existing file
"""

import argparse
import json
import os
import re
import subprocess
import sys
from pathlib import Path

# Constants
KB = 1024
MB = 1024 * KB
GB = 1024 * MB

# Project root
PROJECT_ROOT = Path(__file__).parent.parent

# Default test file location
DEFAULT_TEST_FILE = PROJECT_ROOT / "data" / "100mb.wasot"


def parse_size(size_str: str) -> int:
    """Parse a size string like '100MB' into bytes."""
    size_str = size_str.strip().upper()
    if size_str.endswith("GB"):
        return int(float(size_str[:-2]) * GB)
    elif size_str.endswith("MB"):
        return int(float(size_str[:-2]) * MB)
    elif size_str.endswith("KB"):
        return int(float(size_str[:-2]) * KB)
    elif size_str.endswith("B"):
        return int(size_str[:-1])
    else:
        return int(size_str)


def format_size(bytes_val: int) -> str:
    """Format bytes as human-readable string."""
    if bytes_val >= GB:
        return f"{bytes_val / GB:.2f} GB"
    elif bytes_val >= MB:
        return f"{bytes_val / MB:.2f} MB"
    elif bytes_val >= KB:
        return f"{bytes_val / KB:.2f} KB"
    else:
        return f"{bytes_val} bytes"


def run_command(cmd: list[str], cwd: Path = None, check: bool = True) -> subprocess.CompletedProcess:
    """Run a command and return the result."""
    print(f"  $ {' '.join(cmd)}")
    result = subprocess.run(
        cmd,
        cwd=cwd or PROJECT_ROOT,
        capture_output=True,
        text=True,
    )
    if check and result.returncode != 0:
        print(f"  STDOUT: {result.stdout}")
        print(f"  STDERR: {result.stderr}")
        raise RuntimeError(f"Command failed with exit code {result.returncode}")
    return result


def generate_test_file(output_path: Path, size: int) -> None:
    """Generate a wasot test file of the specified size."""
    print(f"\n📝 Generating {format_size(size)} test file...")

    generator_script = PROJECT_ROOT / "scripts" / "generators" / "generate_large_file.py"

    # Create parent directory if needed
    output_path.parent.mkdir(parents=True, exist_ok=True)

    run_command([
        sys.executable,
        str(generator_script),
        "-o", str(output_path),
        "-s", str(size),
        "--seed", "42",
    ])

    actual_size = output_path.stat().st_size
    print(f"  ✓ Generated: {output_path.name} ({format_size(actual_size)})")


def build_with_profiling() -> Path:
    """Build the release binary with dhat-heap feature."""
    print("\n🔨 Building with memory profiling...")

    run_command([
        "cargo", "build", "--release", "--features", "dhat-heap"
    ])

    binary = PROJECT_ROOT / "target" / "release" / "wasot"
    if not binary.exists():
        raise RuntimeError(f"Binary not found at {binary}")

    print(f"  ✓ Built: {binary}")
    return binary


def run_with_profiler(binary: Path, input_file: Path) -> dict:
    """Run the binary with profiler and return memory stats."""
    print("\n🔬 Running with memory profiler...")

    # dhat outputs to current directory, so we run from project root
    dhat_file = PROJECT_ROOT / "dhat-heap.json"

    # Remove old dhat file if exists
    if dhat_file.exists():
        dhat_file.unlink()

    # Run the command - output to /dev/null since we just care about memory
    result = run_command(
        [str(binary), str(input_file), "-t", "jsonl", "-o", "/dev/null"],
        cwd=PROJECT_ROOT,
        check=True,
    )

    # Parse dhat output from stderr
    stats = parse_dhat_output(result.stderr)

    # Check if JSON was created
    if dhat_file.exists():
        stats["json_path"] = str(dhat_file)

    return stats


def parse_dhat_output(stderr: str) -> dict:
    """Parse dhat summary output from stderr."""
    stats = {}

    # Parse: dhat: At t-gmax: 876,859 bytes in 1,142 blocks
    match = re.search(r"At t-gmax:\s+([\d,]+)\s+bytes", stderr)
    if match:
        stats["peak_bytes"] = int(match.group(1).replace(",", ""))

    # Parse: dhat: At t-end: 68 bytes in 2 blocks
    match = re.search(r"At t-end:\s+([\d,]+)\s+bytes", stderr)
    if match:
        stats["end_bytes"] = int(match.group(1).replace(",", ""))

    # Parse: dhat: Total: 42,626,315,114 bytes in 144,915,589 blocks
    match = re.search(r"Total:\s+([\d,]+)\s+bytes", stderr)
    if match:
        stats["total_bytes"] = int(match.group(1).replace(",", ""))

    return stats


def analyze_dhat_json(json_path: Path) -> list[dict]:
    """Analyze dhat-heap.json for top allocators at peak."""
    with open(json_path) as f:
        data = json.load(f)

    pps = data.get("pps", [])
    ftbl = data.get("ftbl", [])

    # Find top allocators at peak
    top_allocators = sorted(pps, key=lambda x: x.get("gb", 0), reverse=True)[:5]

    results = []
    for pp in top_allocators:
        gb = pp.get("gb", 0)
        fs = pp.get("fs", [])
        frames = [ftbl[f] for f in fs[:3]] if fs else []
        results.append({
            "bytes_at_peak": gb,
            "frames": frames,
        })

    return results


def main():
    parser = argparse.ArgumentParser(
        description="Test that memory usage is bounded during file processing",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python scripts/test_memory_usage.py                    # Default: 100MB file, 5MB limit
  python scripts/test_memory_usage.py --size 200MB       # Larger test file
  python scripts/test_memory_usage.py --max-memory 10MB  # Higher memory limit
  python scripts/test_memory_usage.py --skip-generate    # Reuse existing data/100mb.wasot
  python scripts/test_memory_usage.py -v                 # Show top allocators
        """
    )
    parser.add_argument(
        "--size",
        type=str,
        default="100MB",
        help="Size of test file to generate (default: 100MB)"
    )
    parser.add_argument(
        "--max-memory",
        type=str,
        default="5MB",
        help="Maximum allowed peak memory (default: 5MB)"
    )
    parser.add_argument(
        "--skip-generate",
        action="store_true",
        help="Skip file generation, reuse existing data/100mb.wasot"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show detailed allocation information"
    )
    args = parser.parse_args()

    file_size = parse_size(args.size)
    max_memory = parse_size(args.max_memory)
    test_file = DEFAULT_TEST_FILE

    print("=" * 60)
    print("🧪 Memory Usage Test (DHAT Profiler)")
    print("=" * 60)
    print(f"  Test file:          {test_file}")
    print(f"  Test file size:     {format_size(file_size)}")
    print(f"  Max allowed memory: {format_size(max_memory)}")

    try:
        # Step 1: Generate test file (unless skipped)
        if args.skip_generate:
            if not test_file.exists():
                print(f"\n❌ ERROR: Test file not found: {test_file}")
                print("  Run without --skip-generate to create it")
                return 1
            actual_size = test_file.stat().st_size
            print(f"\n📁 Using existing file: {format_size(actual_size)}")
        else:
            generate_test_file(test_file, file_size)

        # Step 2: Build with profiling
        binary = build_with_profiling()

        # Step 3: Run with profiler
        stats = run_with_profiler(binary, test_file)

        # Step 4: Report results
        print("\n📊 Results")
        print("-" * 40)

        peak_bytes = stats.get("peak_bytes", 0)
        end_bytes = stats.get("end_bytes", 0)
        total_bytes = stats.get("total_bytes", 0)
        input_size = test_file.stat().st_size

        print(f"  Peak memory (t-gmax): {format_size(peak_bytes)}")
        print(f"  End memory (t-end):   {format_size(end_bytes)}")
        print(f"  Total allocated:      {format_size(total_bytes)}")
        print(f"  Input file size:      {format_size(input_size)}")
        print(f"  Memory/Input ratio:   {peak_bytes / input_size:.4f}x")

        # Verbose: show top allocators
        if args.verbose and "json_path" in stats:
            print("\n  Top allocators at peak:")
            allocators = analyze_dhat_json(Path(stats["json_path"]))
            for i, alloc in enumerate(allocators, 1):
                if alloc["bytes_at_peak"] > 0:
                    print(f"    {i}. {format_size(alloc['bytes_at_peak'])}")
                    for frame in alloc["frames"]:
                        # Truncate long frame names
                        if len(frame) > 70:
                            frame = "..." + frame[-67:]
                        print(f"       {frame}")

        # Step 5: Assert memory is bounded
        print("\n" + "=" * 60)

        if peak_bytes <= max_memory:
            print(f"✅ PASSED: Peak memory ({format_size(peak_bytes)}) <= limit ({format_size(max_memory)})")
            return 0
        else:
            print(f"❌ FAILED: Peak memory ({format_size(peak_bytes)}) > limit ({format_size(max_memory)})")
            print(f"\n  Memory is {peak_bytes / max_memory:.1f}x the allowed limit!")
            print(f"  This suggests unbounded memory growth (memory leak).")
            print(f"\n  To diagnose, upload dhat-heap.json to:")
            print(f"  https://nnethercote.github.io/dh_view/dh_view.html")
            return 1

    except Exception as e:
        print(f"\n❌ ERROR: {e}")
        import traceback
        traceback.print_exc()
        return 2


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