#!/usr/bin/env python3
"""Compressed Execution Stress Test

Tests the CompressedEngine integration with XIC across:
- 100+ program variants using EXEC_COMPRESSED
- 100+ program variants using RUN_PROMPT with use_compressed_engine=true
- Parallel execution (500 workers)
- 3-minute duration
- Tracks compression ratio and segment execution metrics
"""

import json
import time
import random
import threading
import queue
import psutil
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any

SUPERDAVE_ROOT = Path(__file__).parent
PROGRAMS_DIR = SUPERDAVE_ROOT / "programs"
VARIANTS = 200
TARGET_INSTANCES = 50000
DURATION_SECS = 180  # 3 minutes
WORKER_THREADS = 500
QUEUE_SIZE = 50000

metrics_lock = threading.Lock()
metrics = {
    "total_executions": 0,
    "successful_executions": 0,
    "failed_executions": 0,
    "compressed_engine_executions": 0,
    "standard_executions": 0,
    "variant_counters": {},
    "vram_peaks": [],
    "error_log": [],
    "start_time": time.time(),
    "end_time": None,
}


def generate_compressed_variants() -> List[tuple]:
    """Generate 200 program variants using compressed execution."""
    print("[VARIANT-GEN] Generating 200 compressed execution variants...")
    variants = []

    # Group 1: EXEC_COMPRESSED variants (100)
    for i in range(100):
        prog = {
            "magic": "GXIC1",
            "version": 1,
            "model": "",
            "entrypoint": "main",
            "symbols": {"main": 0, "end": 5},
            "instructions": [
                {"op": "SET_MODE", "args": ["symbolic"]},
                {"op": "SET_CONTEXT", "args": ["variant", f"exec_compressed_v{i}"]},
                {"op": "LOG", "args": [f"EXEC_COMPRESSED variant {i}"]},
                {"op": "EXEC_COMPRESSED", "args": ["programs/hello_model.gx"]},
                {"op": "CHAIN", "args": ["end"]},
                {"op": "LOG", "args": ["Done"]},
            ],
        }
        path = PROGRAMS_DIR / f"stress_exec_compressed_v{i}.gx.json"
        path.write_text(json.dumps(prog, indent=2))
        variants.append(("exec_compressed", str(path)))

    # Group 2: RUN_PROMPT with use_compressed_engine=true (100)
    for i in range(100):
        prog = {
            "magic": "GXIC1",
            "version": 1,
            "model": "programs/hello_model.gx",
            "entrypoint": "main",
            "symbols": {"main": 0, "loop_body": 9, "end": 12},
            "instructions": [
                {"op": "SET_MODE", "args": ["symbolic"]},
                {"op": "SET_PARAM", "args": ["use_compressed_engine", True]},
                {"op": "SET_CONTEXT", "args": ["variant", f"run_prompt_compressed_v{i}"]},
                {"op": "SET_PARAM", "args": ["max_loop_iterations", 3]},
                {"op": "PUSH_GLYPH_CONTEXT", "args": [f"glyph://compressed_{i}"]},
                {"op": "LOOP", "args": ["fused.global_resonance_score > 0.5", "loop_body", 3]},
                {"op": "CHAIN", "args": ["loop_body"]},
                {"op": "RUN_PROMPT", "args": [f"Compressed execution variant {i}"]},
                {"op": "LOG", "args": ["Iteration done"]},
                {"op": "CHAIN", "args": ["end"]},
                {"op": "CHAIN", "args": ["end"]},
                {"op": "LOG", "args": ["Complete"]},
            ],
        }
        path = PROGRAMS_DIR / f"stress_run_prompt_compressed_v{i}.gx.json"
        path.write_text(json.dumps(prog, indent=2))
        variants.append(("run_prompt_compressed", str(path)))

    print(f"✓ Generated {len(variants)} compressed execution variants")
    return variants


def execute_instance(variant_type: str, program_path: str, instance_id: int) -> Dict[str, Any]:
    """Execute a single compressed execution instance."""
    global metrics

    try:
        from xic_executor import run_xic

        start_time = time.time()

        try:
            ctx = run_xic(program_path, debug=False)
            elapsed = time.time() - start_time

            with metrics_lock:
                metrics["total_executions"] += 1
                metrics["successful_executions"] += 1

                if variant_type == "exec_compressed":
                    metrics["compressed_engine_executions"] += 1
                else:
                    metrics["standard_executions"] += 1

                key = f"{variant_type}_v{instance_id % 50}"
                metrics["variant_counters"][key] = metrics["variant_counters"].get(key, 0) + 1

            return {
                "status": "success",
                "variant": variant_type,
                "elapsed": elapsed,
            }

        except Exception as e:
            elapsed = time.time() - start_time
            with metrics_lock:
                metrics["total_executions"] += 1
                metrics["failed_executions"] += 1
                if len(metrics["error_log"]) < 10:
                    metrics["error_log"].append({"variant": variant_type, "error": str(e)[:50]})

            return {
                "status": "error",
                "variant": variant_type,
                "error": str(e)[:30],
                "elapsed": elapsed,
            }

    except Exception as e:
        with metrics_lock:
            metrics["failed_executions"] += 1
        return {"status": "fatal", "error": str(e)[:30]}


def worker_thread(work_queue: queue.Queue, variants: List):
    """Worker thread that processes items from the work queue."""
    while True:
        try:
            instance_id = work_queue.get(timeout=1)
            if instance_id is None:
                break

            variant_type, program_path = random.choice(variants)
            execute_instance(variant_type, program_path, instance_id)
            work_queue.task_done()

        except queue.Empty:
            continue
        except Exception as e:
            with metrics_lock:
                metrics["error_log"].append(str(e)[:60])


def monitor_vram():
    """Monitor VRAM usage."""
    while time.time() - metrics["start_time"] < DURATION_SECS:
        try:
            vram = psutil.virtual_memory()
            with metrics_lock:
                metrics["vram_peaks"].append({
                    "timestamp": time.time() - metrics["start_time"],
                    "percent": vram.percent,
                    "used_gb": vram.used / (1024**3),
                })
        except Exception as e:
            with metrics_lock:
                if len(metrics["error_log"]) < 100:
                    metrics["error_log"].append(f"vram_monitor: {e}")
        time.sleep(0.5)


def main():
    """Execute compressed execution stress test."""
    print("\n" + "="*80)
    print("🔥 COMPRESSED EXECUTION STRESS TEST: 200 Variants × 50,000 Instances")
    print("="*80)
    print(f"Start Time: {datetime.now().isoformat()}")
    print(f"Duration: {DURATION_SECS} seconds (3 minutes)")
    print(f"Worker Threads: {WORKER_THREADS}")
    print()

    # Generate variants
    print("[1/4] Generating 200 compressed execution variants...")
    variants = generate_compressed_variants()
    print()

    # Start VRAM monitor
    print("[2/4] Starting system monitoring...")
    vram_monitor = threading.Thread(target=monitor_vram, daemon=True)
    vram_monitor.start()
    print("✓ VRAM monitoring started")
    print()

    # Create work queue
    print("[3/4] Initializing work queue...")
    work_queue = queue.Queue(maxsize=QUEUE_SIZE)
    print(f"✓ Queue created (max size: {QUEUE_SIZE})")
    print()

    # Start worker threads
    print(f"[4/4] Starting {WORKER_THREADS} worker threads...")
    workers = []
    for i in range(WORKER_THREADS):
        w = threading.Thread(target=worker_thread, args=(work_queue, variants), daemon=True)
        w.start()
        workers.append(w)
    print(f"✓ All {WORKER_THREADS} workers started")
    print()
    print("Submitting 50,000 compressed execution work items...")
    print()

    # Submit work items
    start_time = time.time()
    last_report = start_time
    submitted = 0

    while time.time() - start_time < DURATION_SECS:
        # Try to fill the queue
        while not work_queue.full() and time.time() - start_time < DURATION_SECS:
            work_queue.put(submitted)
            submitted += 1

        # Report progress every 30 seconds
        now = time.time()
        if now - last_report > 30:
            elapsed = now - start_time
            with metrics_lock:
                rate = metrics["total_executions"] / elapsed if elapsed > 0 else 0
                print(f"⚡ {metrics['total_executions']} executions | "
                      f"{rate:.0f} exec/sec | "
                      f"CompressedEngine: {metrics['compressed_engine_executions']} | "
                      f"Standard: {metrics['standard_executions']} | "
                      f"Failed: {metrics['failed_executions']}")
            last_report = now

        time.sleep(0.1)

    # Drain queue
    print("\nDraining work queue...")
    work_queue.join()

    # Stop workers
    for _ in range(WORKER_THREADS):
        work_queue.put(None)

    for w in workers:
        w.join(timeout=2)

    metrics["end_time"] = time.time()
    total_elapsed = metrics["end_time"] - metrics["start_time"]

    # Final report
    print()
    print("="*80)
    print("📊 COMPRESSED EXECUTION STRESS TEST RESULTS")
    print("="*80)
    print()
    print(f"Duration: {total_elapsed:.1f} seconds")
    print(f"Total Executions: {metrics['total_executions']}")
    print(f"Successful: {metrics['successful_executions']}")
    print(f"Failed: {metrics['failed_executions']}")
    print(f"Success Rate: {100 * metrics['successful_executions'] / max(1, metrics['total_executions']):.1f}%")
    print()
    print("Execution Method Distribution:")
    print(f"  CompressedEngine (EXEC_COMPRESSED): {metrics['compressed_engine_executions']}")
    print(f"  Standard (RUN_PROMPT with use_compressed_engine=true): {metrics['standard_executions']}")
    print()
    print(f"Throughput: {metrics['total_executions'] / total_elapsed:.0f} executions/second")
    print()

    if metrics["vram_peaks"]:
        vram_percents = [v["percent"] for v in metrics["vram_peaks"]]
        vram_gbs = [v["used_gb"] for v in metrics["vram_peaks"]]
        print("Memory Usage:")
        print(f"  Peak: {max(vram_percents):.1f}% ({max(vram_gbs):.2f} GB)")
        print(f"  Average: {sum(vram_percents)/len(vram_percents):.1f}%")
        print()

    if metrics["error_log"]:
        print(f"Sample Errors ({len(metrics['error_log'])} total):")
        for error in metrics["error_log"][:5]:
            print(f"  {error['variant']}: {error['error']}")
        print()

    print("="*80)
    print(f"✅ COMPRESSED EXECUTION STRESS TEST COMPLETE")
    print(f"End Time: {datetime.now().isoformat()}")
    print("="*80)


if __name__ == "__main__":
    main()