#!/usr/bin/env python3
"""
Deep diagnostic tool to understand what causes <8s gap failures.
Monitors: GPU memory, timing, queue states, torch internals.
"""

import torch
import time
import requests
import json
import subprocess
import threading
from collections import defaultdict

def get_gpu_memory():
    """Get current GPU memory usage."""
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3  # GB
        reserved = torch.cuda.memory_reserved() / 1024**3    # GB
        return allocated, reserved
    return 0, 0

def get_gpu_clock_speeds():
    """Get GPU clock speeds via rocm-smi."""
    try:
        result = subprocess.run(['rocm-smi', '--showclocks'], 
                              capture_output=True, text=True, timeout=1)
        lines = result.stdout.split('\n')
        sclk = mclk = 0
        for line in lines:
            if 'sclk' in line.lower() and 'level' in line.lower():
                # Parse "level: X: (YYYYMhz)"
                parts = line.split('(')
                if len(parts) > 1:
                    mhz = parts[1].split('Mhz')[0]
                    if mhz.isdigit():
                        sclk = int(mhz)
            elif 'mclk' in line.lower() and 'level' in line.lower():
                parts = line.split('(')
                if len(parts) > 1:
                    mhz = parts[1].split('Mhz')[0]
                    if mhz.isdigit():
                        mclk = int(mhz)
        return sclk, mclk
    except:
        return 0, 0

def monitor_job_execution(job_id, log_file):
    """Monitor a single job execution in detail."""
    start_time = time.time()
    last_check = start_time
    
    print(f"\n{'='*60}")
    print(f"Monitoring Job: {job_id}")
    print(f"{'='*60}")
    
    with open(log_file, 'a') as f:
        f.write(f"\n{'='*60}\n")
        f.write(f"Job: {job_id} | Started: {time.strftime('%H:%M:%S')}\n")
        f.write(f"{'='*60}\n")
        
        # Initial state
        mem_alloc, mem_reserved = get_gpu_memory()
        sclk, mclk = get_gpu_clock_speeds()
        f.write(f"INITIAL STATE:\n")
        f.write(f"  GPU Mem: {mem_alloc:.2f}GB allocated, {mem_reserved:.2f}GB reserved\n")
        f.write(f"  GPU Clocks: sclk={sclk}MHz, mclk={mclk}MHz\n")
        f.write(f"  PyTorch cache: {torch.cuda.memory_reserved()/1024**3:.2f}GB\n\n")
        
        # Monitor during execution (poll every 0.5s for ~10s)
        for i in range(20):
            time.sleep(0.5)
            elapsed = time.time() - start_time
            
            mem_alloc, mem_reserved = get_gpu_memory()
            sclk, mclk = get_gpu_clock_speeds()
            
            f.write(f"T+{elapsed:.1f}s: mem={mem_alloc:.2f}GB, "
                   f"sclk={sclk}MHz, mclk={mclk}MHz\n")
            
            # If job seems done (very rough heuristic), break
            if elapsed > 8 and mem_alloc < 0.5:
                break
        
        # Final state
        time.sleep(1)
        elapsed = time.time() - start_time
        mem_alloc, mem_reserved = get_gpu_memory()
        sclk, mclk = get_gpu_clock_speeds()
        
        f.write(f"\nFINAL STATE (T+{elapsed:.1f}s):\n")
        f.write(f"  GPU Mem: {mem_alloc:.2f}GB allocated, {mem_reserved:.2f}GB reserved\n")
        f.write(f"  GPU Clocks: sclk={sclk}MHz, mclk={mclk}MHz\n")
        f.write(f"  PyTorch cache: {torch.cuda.memory_reserved()/1024**3:.2f}GB\n")
        f.write(f"{'='*60}\n\n")
        
    print(f"Job {job_id} monitoring complete (T+{elapsed:.1f}s)")

def test_gap_with_monitoring(gap_seconds, num_jobs=3):
    """Test a specific gap with deep monitoring."""
    log_file = f'/tmp/gpu_diagnostic_{gap_seconds}s.log'
    
    print(f"\n{'#'*60}")
    print(f"# Testing {gap_seconds}s gap with {num_jobs} jobs")
    print(f"# Log file: {log_file}")
    print(f"{'#'*60}")
    
    # Clear log
    with open(log_file, 'w') as f:
        f.write(f"GPU Diagnostic Test - {gap_seconds}s gap\n")
        f.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"{'='*60}\n\n")
    
    for i in range(num_jobs):
        print(f"\n[Job {i+1}/{num_jobs}]")
        
        # Submit job
        text = f"Diagnostic test job {i+1} with {gap_seconds} second gap for deep analysis."
        response = requests.post(
            'http://localhost:8765/api/speak',
            headers={'Content-Type': 'application/json'},
            json={'text': text}
        )
        
        if response.status_code == 200:
            data = response.json()
            job_id = data.get('job_id', 'unknown')[:8]
            print(f"  Job ID: {job_id}")
            print(f"  Queue size: {data.get('queue_size', '?')}")
            
            # Monitor this job
            monitor_thread = threading.Thread(
                target=monitor_job_execution,
                args=(job_id, log_file)
            )
            monitor_thread.start()
            
            # Wait for gap
            print(f"  Waiting {gap_seconds}s before next job...")
            time.sleep(gap_seconds)
            
            monitor_thread.join(timeout=1)  # Don't block too long
        else:
            print(f"  ERROR: {response.status_code}")
    
    # Wait for last job
    print(f"\nWaiting 10s for last job to complete...")
    time.sleep(10)
    
    print(f"\nDiagnostic log saved to: {log_file}")
    print(f"Timing results in: /tmp/soprano_timing_debug.txt")

def compare_torch_state():
    """Compare PyTorch internal state between good and bad runs."""
    print("\n=== PyTorch State Analysis ===")
    print(f"CUDA available: {torch.cuda.is_available()}")
    print(f"CUDA device: {torch.cuda.get_device_name(0)}")
    print(f"Current stream: {torch.cuda.current_stream()}")
    print(f"Memory allocated: {torch.cuda.memory_allocated()/1024**3:.2f}GB")
    print(f"Memory reserved: {torch.cuda.memory_reserved()/1024**3:.2f}GB")
    print(f"Max memory allocated: {torch.cuda.max_memory_allocated()/1024**3:.2f}GB")
    print(f"Max memory reserved: {torch.cuda.max_memory_reserved()/1024**3:.2f}GB")
    print(f"Memory cached: {torch.cuda.memory_cached()/1024**3:.2f}GB" if hasattr(torch.cuda, 'memory_cached') else "N/A")

if __name__ == '__main__':
    print("="*60)
    print("GPU Diagnostic Tool for Job Spacing Investigation")
    print("="*60)
    
    # Clear timing log
    import os
    if os.path.exists('/tmp/soprano_timing_debug.txt'):
        os.remove('/tmp/soprano_timing_debug.txt')
    
    # Show initial state
    compare_torch_state()
    
    # Test 8s gap (should be good)
    print("\n" + "="*60)
    print("TEST 1: 8s gap (expected: GOOD)")
    print("="*60)
    test_gap_with_monitoring(8, num_jobs=3)
    
    print("\n" + "="*60)
    print("Waiting 15s to let GPU fully settle...")
    print("="*60)
    time.sleep(15)
    
    # Test 7s gap (should be bad)
    print("\n" + "="*60)
    print("TEST 2: 7s gap (expected: BAD)")
    print("="*60)
    test_gap_with_monitoring(7, num_jobs=3)
    
    print("\n" + "="*60)
    print("Analysis Complete!")
    print("="*60)
    
    # Show results
    print("\nPerformance Results:")
    with open('/tmp/soprano_timing_debug.txt', 'r') as f:
        lines = f.readlines()
        for i, line in enumerate(lines):
            if 'Realtime' in line:
                print(f"  Job {(i//6)+1}: {line.strip()}")
                if i+1 < len(lines):
                    print(f"           {lines[i+1].strip()}")
    
    print("\nCheck detailed logs:")
    print("  - /tmp/gpu_diagnostic_8s.log")
    print("  - /tmp/gpu_diagnostic_7s.log")
