import subprocess
import sys
import os
import json
import logging

logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')

def get_python_exe():
    """Detects the best python executable to use (prefers venv)."""
    # Prefer explicit venv path
    venv_python = os.path.join(os.getcwd(), "venv/bin/python3")
    if os.path.exists(venv_python):
        return venv_python
    return sys.executable

def run_benchmark():
    python_exe = get_python_exe()
    logging.info(f"Using Python executable: {python_exe}")
    
    modes = ["fp32", "int8", "selective"]
    summary = {}
    
    os.makedirs("results", exist_ok=True)
    os.makedirs("outputs", exist_ok=True)
    
    for mode in modes:
        logging.info(f"=== BENCHMARKING MODE: {mode} ===")
        output_file = f"results/results_{mode}.json"
        
        # Run in subprocess to ensure isolated memory measurement
        cmd = [
            python_exe, "src/optimize_tts.py",
            "--mode", mode,
            "--output_json", output_file,
            "--text", "This is a benchmarking sample for CPU optimized MOSS TTS. It tests end-to-end latency."
        ]
        
        try:
            # Capturing outputs
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode == 0:
                if os.path.exists(output_file):
                    with open(output_file, "r") as f:
                        summary[mode] = json.load(f)
                    logging.info(f"Success: {mode}")
                else:
                    logging.error(f"Output file {output_file} not found for mode {mode}")
            else:
                logging.error(f"Failed to benchmark {mode}. Return code: {result.returncode}")
                logging.error(f"STDERR: {result.stderr}")
        except Exception as e:
            logging.error(f"Subprocess error for {mode}: {e}")

    # Print Summary Table
    print("\n" + "="*80)
    print(f"{'Quantization Mode':<20} | {'RAM (MB)':<12} | {'Latency (ms)':<15} | {'Load (s)':<10}")
    print("-" * 80)
    for mode in modes:
        if mode in summary:
            data = summary[mode]
            print(f"{mode:<20} | {data['peak_ram_mb']:<12.2f} | {data['latency_ms']:<15.2f} | {data['load_time_sec']:<10.2f}")
        else:
            print(f"{mode:<20} | {'FAILED':<12} | {'N/A':<15} | {'N/A':<10}")
    print("="*80 + "\n")

    with open("results/benchmark_summary.json", "w") as f:
        json.dump(summary, f, indent=4)
    logging.info("Benchmark summary saved to results/benchmark_summary.json")

if __name__ == "__main__":
    run_benchmark()