#!/usr/bin/env python3
"""
Profiling helper script for wasot.

Provides easy access to DHAT heap profiling and analysis.

Usage:
    ./scripts/profile.py                    # Run DHAT profile on 1mb.wasot
    ./scripts/profile.py --file FILE        # Profile specific file
    ./scripts/profile.py --analyze          # Analyze existing dhat-heap.json
    ./scripts/profile.py --compare A B      # Compare two DHAT profiles
"""

import argparse
import json
import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path

SCRIPT_DIR = Path(__file__).parent.resolve()
PROJECT_ROOT = SCRIPT_DIR.parent
VENV_PYTHON = PROJECT_ROOT / ".venv" / "bin" / "python"


def ensure_venv():
    """Re-exec with venv Python if not already running in it."""
    if not VENV_PYTHON.exists():
        return  # No venv, continue with system Python

    current_python = Path(sys.executable).resolve()
    if current_python != VENV_PYTHON.resolve():
        os.execv(str(VENV_PYTHON), [str(VENV_PYTHON)] + sys.argv)


ensure_venv()


@dataclass
class DhatSummary:
    """Summary of DHAT profile."""
    total_bytes: int
    total_allocs: int
    peak_bytes: int
    peak_blocks: int
    top_sites: list


def build_with_dhat() -> Path:
    """Build wasot with DHAT support."""
    print("Building wasot with DHAT support...")

    env = os.environ.copy()
    env["CARGO_PROFILE_RELEASE_DEBUG"] = "true"
    env["CARGO_PROFILE_RELEASE_STRIP"] = "false"

    result = subprocess.run(
        ["cargo", "build", "--release", "--bin", "wasot", "--features", "dhat-heap"],
        cwd=PROJECT_ROOT,
        capture_output=True,
        text=True,
        env=env,
    )

    if result.returncode != 0:
        print("Build failed:", file=sys.stderr)
        print(result.stderr, file=sys.stderr)
        sys.exit(1)

    print("Build complete.")
    return PROJECT_ROOT / "target" / "release" / "wasot"


def run_profile(binary: Path, input_file: Path) -> Path:
    """Run the binary and generate DHAT profile."""
    print(f"Profiling: {input_file}")

    result = subprocess.run(
        [str(binary), str(input_file), "-t", "wasot"],
        capture_output=True,
        text=True,
        cwd=PROJECT_ROOT,
    )

    # DHAT output goes to stderr
    for line in result.stderr.split('\n'):
        if line.startswith('dhat:'):
            print(f"  {line}")

    dhat_file = PROJECT_ROOT / "dhat-heap.json"
    if not dhat_file.exists():
        print("Error: dhat-heap.json not generated", file=sys.stderr)
        sys.exit(1)

    return dhat_file


def analyze_dhat(dhat_file: Path, top_n: int = 15) -> DhatSummary:
    """Analyze DHAT profile and return summary."""
    with open(dhat_file) as f:
        d = json.load(f)

    ftbl = d.get('ftbl', [])
    pps = d.get('pps', [])

    total_bytes = sum(pp.get('tb', 0) for pp in pps)
    total_allocs = sum(pp.get('tblk', 0) for pp in pps)
    peak_bytes = d.get('tGmax', 0)
    peak_blocks = d.get('tGmaxBlk', 0)

    # Get top allocation sites
    sorted_pps = sorted(pps, key=lambda x: x.get('tb', 0), reverse=True)[:top_n]

    top_sites = []
    for pp in sorted_pps:
        tb = pp.get('tb', 0)
        tblk = pp.get('tblk', 0)
        mb = pp.get('mb', 0)
        fs = pp.get('fs', [])

        # Get meaningful frames (skip alloc internals)
        frames = []
        for frame_idx in fs[:8]:
            if frame_idx < len(ftbl):
                frame = ftbl[frame_idx]
                # Skip low-level allocator frames
                if 'exchange_malloc' in frame or 'Allocator>::allocate' in frame:
                    continue
                frames.append(frame)
                if len(frames) >= 3:
                    break

        top_sites.append({
            'bytes': tb,
            'allocs': tblk,
            'max_live': mb,
            'frames': frames,
        })

    return DhatSummary(
        total_bytes=total_bytes,
        total_allocs=total_allocs,
        peak_bytes=peak_bytes,
        peak_blocks=peak_blocks,
        top_sites=top_sites,
    )


def format_bytes(b: int) -> str:
    """Format bytes as human-readable."""
    if b >= 1_000_000_000:
        return f"{b / 1_000_000_000:.1f} GB"
    elif b >= 1_000_000:
        return f"{b / 1_000_000:.1f} MB"
    elif b >= 1_000:
        return f"{b / 1_000:.1f} KB"
    else:
        return f"{b} B"


def print_summary(summary: DhatSummary, title: str = "DHAT Profile Summary"):
    """Print formatted summary."""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)
    print()
    print(f"  Total allocations:  {format_bytes(summary.total_bytes):>12} ({summary.total_allocs:,} blocks)")
    print(f"  Peak heap usage:    {format_bytes(summary.peak_bytes):>12} ({summary.peak_blocks:,} blocks)")
    print()
    print("-" * 70)
    print("  TOP ALLOCATION SITES")
    print("-" * 70)

    for i, site in enumerate(summary.top_sites[:10]):
        print(f"\n  #{i+1}: {format_bytes(site['bytes']):>10} ({site['allocs']:,} allocs, max live: {format_bytes(site['max_live'])})")
        for frame in site['frames'][:2]:
            # Clean up frame for display
            if len(frame) > 80:
                # Extract the function name part
                if ': ' in frame:
                    frame = frame.split(': ', 1)[1]
                if len(frame) > 80:
                    frame = frame[:77] + '...'
            print(f"       {frame}")

    print()
    print("-" * 70)

    # Check for common issues
    print("\n  ANALYSIS")
    print("-" * 70)

    issues = []

    # Check for regex allocations
    regex_bytes = 0
    for site in summary.top_sites:
        for frame in site['frames']:
            if 'regex' in frame.lower():
                regex_bytes += site['bytes']
                break

    if regex_bytes > 1_000_000:
        issues.append(f"  WARNING: {format_bytes(regex_bytes)} in regex operations - check for Regex::new in loops")

    # Check total allocation vs reasonable
    if summary.total_bytes > 100_000_000:
        issues.append(f"  WARNING: High total allocations ({format_bytes(summary.total_bytes)}) - look for optimization opportunities")

    if not issues:
        issues.append("  OK: No obvious allocation issues detected")

    for issue in issues:
        print(issue)

    print()


def compare_profiles(file_a: Path, file_b: Path):
    """Compare two DHAT profiles."""
    summary_a = analyze_dhat(file_a)
    summary_b = analyze_dhat(file_b)

    print()
    print("=" * 70)
    print("  PROFILE COMPARISON")
    print("=" * 70)
    print()
    print(f"  Profile A: {file_a.name}")
    print(f"  Profile B: {file_b.name}")
    print()
    print("-" * 70)
    print(f"  {'Metric':<25} {'Profile A':>15} {'Profile B':>15} {'Change':>12}")
    print("-" * 70)

    def fmt_change(a: int, b: int) -> str:
        if a == 0:
            return "N/A"
        pct = ((b - a) / a) * 100
        if pct > 0:
            return f"+{pct:.1f}%"
        else:
            return f"{pct:.1f}%"

    print(f"  {'Total bytes':<25} {format_bytes(summary_a.total_bytes):>15} {format_bytes(summary_b.total_bytes):>15} {fmt_change(summary_a.total_bytes, summary_b.total_bytes):>12}")
    print(f"  {'Total allocs':<25} {summary_a.total_allocs:>15,} {summary_b.total_allocs:>15,} {fmt_change(summary_a.total_allocs, summary_b.total_allocs):>12}")
    print(f"  {'Peak bytes':<25} {format_bytes(summary_a.peak_bytes):>15} {format_bytes(summary_b.peak_bytes):>15} {fmt_change(summary_a.peak_bytes, summary_b.peak_bytes):>12}")
    print("-" * 70)
    print()


def main():
    parser = argparse.ArgumentParser(
        description="Profiling helper for wasot",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    ./scripts/profile.py                     # Profile with 1mb.wasot
    ./scripts/profile.py --file my.wasot     # Profile specific file
    ./scripts/profile.py --analyze           # Analyze existing dhat-heap.json
    ./scripts/profile.py --compare a.json b.json  # Compare two profiles
        """
    )
    parser.add_argument(
        "--file", "-f",
        type=Path,
        default=PROJECT_ROOT / "data" / "bench" / "1mb.wasot",
        help="Input file to profile (default: data/bench/1mb.wasot)",
    )
    parser.add_argument(
        "--analyze", "-a",
        action="store_true",
        help="Only analyze existing dhat-heap.json (don't run new profile)",
    )
    parser.add_argument(
        "--compare", "-c",
        nargs=2,
        type=Path,
        metavar=("A", "B"),
        help="Compare two DHAT profiles",
    )
    parser.add_argument(
        "--output", "-o",
        type=Path,
        help="Save DHAT profile to specified location",
    )
    args = parser.parse_args()

    if args.compare:
        compare_profiles(args.compare[0], args.compare[1])
        return

    if args.analyze:
        dhat_file = PROJECT_ROOT / "dhat-heap.json"
        if not dhat_file.exists():
            print("Error: dhat-heap.json not found. Run without --analyze first.", file=sys.stderr)
            sys.exit(1)
        summary = analyze_dhat(dhat_file)
        print_summary(summary)
        return

    # Build and run profile
    binary = build_with_dhat()
    dhat_file = run_profile(binary, args.file)

    # Save profile if requested
    if args.output:
        import shutil
        shutil.copy(dhat_file, args.output)
        print(f"Profile saved to: {args.output}")

    # Analyze and print
    summary = analyze_dhat(dhat_file)
    print_summary(summary, f"DHAT Profile: {args.file.name}")


if __name__ == "__main__":
    main()
