"""
Benchmark engine for testing LLM models for Zelin.
"""
import os
import sys
import json
import time
import gc
import subprocess
from pathlib import Path

# Install dependencies at runtime with verbose output
print("=== Installing dependencies ===", flush=True)
print("Installing psutil and huggingface_hub...", flush=True)
result = subprocess.run([sys.executable, "-m", "pip", "install", "psutil", "huggingface_hub"], capture_output=True, text=True)
print(f"  Return code: {result.returncode}", flush=True)
if result.returncode != 0:
    print(f"  STDERR: {result.stderr[-1000:]}", flush=True)

print("Installing llama-cpp-python (pre-built wheel)...", flush=True)
result = subprocess.run(
    [sys.executable, "-m", "pip", "install", "llama-cpp-python", "--extra-index-url", "https://abetlen.github.io/llama-cpp-python/whl/cpu"],
    capture_output=True, text=True
)
print(f"  Return code: {result.returncode}", flush=True)
print(f"  STDOUT (last 1000): {result.stdout[-1000:]}", flush=True)
if result.returncode != 0:
    print(f"  STDERR: {result.stderr[-1000:]}", flush=True)
    print("Trying install from source...", flush=True)
    result2 = subprocess.run([sys.executable, "-m", "pip", "install", "llama-cpp-python"], capture_output=True, text=True)
    print(f"  Return code: {result2.returncode}", flush=True)
    print(f"  STDOUT (last 1000): {result2.stdout[-1000:]}", flush=True)

# Verify - add user site-packages to path first
import site
import os
user_site = site.getusersitepackages()
print(f"User site-packages: {user_site}", flush=True)
if user_site not in sys.path:
    sys.path.insert(0, user_site)
    print(f"Added to sys.path", flush=True)

# Also try adding common locations
for p in [
    os.path.expanduser('~/.local/lib/python3.11/site-packages'),
    os.path.expanduser('~/.local/lib/python3.10/site-packages'),
    os.path.expanduser('~/.local/lib/python3.12/site-packages'),
]:
    if os.path.exists(p) and p not in sys.path:
        sys.path.insert(0, p)
        print(f"Also added: {p}", flush=True)

print(f"sys.path: {sys.path[:5]}", flush=True)

print("Verifying llama_cpp import...", flush=True)
try:
    import llama_cpp
    print(f"✅ llama_cpp imported!", flush=True)
except ImportError as e:
    print(f"❌ llama_cpp import failed: {e}", flush=True)
    # List what's in user site-packages
    if os.path.exists(user_site):
        print(f"Contents of {user_site}:", flush=True)
        for f in os.listdir(user_site)[:20]:
            print(f"  {f}", flush=True)
    print("ABORTING - cannot run benchmark without llama_cpp", flush=True)
    sys.exit(1)

import psutil
print("All deps ready!", flush=True)

# Models to benchmark (in order from smallest to biggest)
MODELS = [
    # SKIP: TeapotLLM has a known bug with llama-cpp-python (GGML_ASSERT cross-attention failure)
    # {
    #     "name": "TeapotLLM",
    #     "repo": "mradermacher/teapotllm-chat-GGUF",
    #     "file": "teapotllm-chat.Q4_K_M.gguf",
    #     "size_mb": 464,
    #     "params": "0.8B",
    #     "context": 512,
    #     "language": "en",
    # },
    {
        "name": "TinyLlama-1.1B",
        "repo": "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
        "file": "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
        "size_mb": 638,
        "params": "1.1B",
        "context": 2048,
        "language": "en",
    },
    {
        "name": "SmolLM2-1.7B",
        "repo": "unsloth/SmolLM2-1.7B-Instruct-GGUF",
        "file": "SmolLM2-1.7B-Instruct-Q4_K_M.gguf",
        "size_mb": 1007,
        "params": "1.7B",
        "context": 8192,
        "language": "multi",
    },
    {
        "name": "Phi-4-mini",
        "repo": "unsloth/Phi-4-mini-instruct-GGUF",
        "file": "Phi-4-mini-instruct-Q4_K_M.gguf",
        "size_mb": 2376,
        "params": "3.8B",
        "context": 128000,
        "language": "multi",
    },
    {
        "name": "Mistral-7B",
        "repo": "TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
        "file": "mistral-7b-instruct-v0.2.Q4_K_M.gguf",
        "size_mb": 4166,
        "params": "7B",
        "context": 32768,
        "language": "multi",
    },
    {
        "name": "Salamandra-7B",
        "repo": "cstr/salamandra-7b-instruct-GGUF",
        "file": "salamandra-7b-instruct.Q4_K_M-f32.gguf",
        "size_mb": 4626,
        "params": "7B",
        "context": 8192,
        "language": "es",
    },
    {
        "name": "Llama-3.1-8B",
        "repo": "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
        "file": "Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
        "size_mb": 4693,
        "params": "8B",
        "context": 128000,
        "language": "multi",
    },
    {
        "name": "Aya-23-8B",
        "repo": "bartowski/aya-23-8B-GGUF",
        "file": "aya-23-8B-Q4_K_M.gguf",
        "size_mb": 4823,
        "params": "8B",
        "context": 8192,
        "language": "multi",
    },
    {
        "name": "Qwen3-8B",
        "repo": "unsloth/Qwen3-8B-GGUF",
        "file": "Qwen3-8B-Q4_K_M.gguf",
        "size_mb": 4795,
        "params": "8B",
        "context": 128000,
        "language": "multi",
    },
    {
        "name": "Gemma-2-9B",
        "repo": "bartowski/gemma-2-9b-it-GGUF",
        "file": "gemma-2-9b-it-Q4_K_M-fp16.gguf",
        "size_mb": 6526,
        "params": "9B",
        "context": 8192,
        "language": "multi",
    },
]


# Benchmark prompts (in Spanish, for Zelin use cases)
BENCHMARK_PROMPTS = [
    {
        "id": "greeting_casual",
        "prompt": "Hola, ¿qué tal? ¿Cómo estás?",
        "max_tokens": 150,
        "category": "social",
        "description": "Saludo casual - debe responder como mexicana, no argentina",
    },
    {
        "id": "minecraft_help",
        "prompt": "¿Cómo hago una espada de diamante en Minecraft? Dame los pasos exactos.",
        "max_tokens": 300,
        "category": "minecraft",
        "description": "Conocimiento de Minecraft - debe ser preciso",
    },
    {
        "id": "moderation",
        "prompt": "Un usuario del servidor dijo: 'Oye tú, vete de aquí, eres un inútil'. ¿Qué deberías hacer como moderadora?",
        "max_tokens": 250,
        "category": "moderation",
        "description": "Capacidad de moderación Discord",
    },
    {
        "id": "reasoning",
        "prompt": "Si tengo 3 manzanas y le doy una a mi hermana, y luego compro 5 más, ¿cuántas manzanas tengo? Explica tu razonamiento.",
        "max_tokens": 200,
        "category": "reasoning",
        "description": "Razonamiento matemático básico",
    },
    {
        "id": "roleplay",
        "prompt": "Eres Zelin, una chica mexicana jugadora de Minecraft. Acabas de morir en el juego por un skeleton. Reacciona en menos de 50 palabras.",
        "max_tokens": 100,
        "category": "roleplay",
        "description": "Roleplay con personalidad mexicana",
    },
    {
        "id": "knowledge",
        "prompt": "¿Cuál es la capital de Australia? ¿Y la de Canadá? Responde solo las capitales.",
        "max_tokens": 50,
        "category": "knowledge",
        "description": "Conocimiento general - debe ser preciso",
    },
    {
        "id": "emoji",
        "prompt": "Saluda al servidor con emojis de Discord. Usa el formato :nombre_emoji: nada más.",
        "max_tokens": 50,
        "category": "format",
        "description": "Formato correcto de emojis Discord",
    },
]


def download_model(model_info, model_dir):
    """Download the GGUF model file from HF Hub."""
    from huggingface_hub import hf_hub_download
    print(f"  [download] {model_info['name']} from {model_info['repo']}")
    print(f"  [download] file: {model_info['file']} ({model_info['size_mb']} MB)")
    
    start = time.time()
    local_path = hf_hub_download(
        repo_id=model_info['repo'],
        filename=model_info['file'],
        local_dir=model_dir,
        cache_dir="/tmp/hf_cache",
    )
    elapsed = time.time() - start
    print(f"  [download] done in {elapsed:.1f}s, path: {local_path}")
    return local_path, elapsed


def run_model_benchmark(model_path, model_info):
    """Run llama-cpp-python benchmark on the model."""
    from llama_cpp import Llama
    
    print(f"  [load] loading model...")
    load_start = time.time()
    
    # Determine n_ctx based on model size
    # Cap at model's training context AND 4096 (whichever is smaller)
    # Some models (TeapotLLM) were trained with only 512 context
    n_ctx = min(model_info["context"], 4096)  # Cap at 4K for benchmark
    # Some models have very small training context, use that
    if model_info["context"] <= 1024:
        n_ctx = model_info["context"]
    n_threads = min(4, os.cpu_count() or 2)
    
    try:
        llm = Llama(
            model_path=model_path,
            n_ctx=n_ctx,
            n_threads=n_threads,
            n_gpu_layers=0,  # CPU only
            verbose=False,
            use_mlock=False,
            use_mmap=True,
        )
    except Exception as e:
        print(f"  [load] FAILED: {e}")
        return None
    
    load_time = time.time() - load_start
    print(f"  [load] done in {load_time:.1f}s")
    
    # Check memory usage
    process = psutil.Process()
    mem_after_load = process.memory_info().rss / (1024 * 1024)  # MB
    
    results = {
        "model_info": model_info,
        "load_time_s": round(load_time, 2),
        "memory_after_load_mb": round(mem_after_load, 1),
        "prompts": [],
    }
    
    # Run each prompt
    total_tokens = 0
    total_time = 0
    
    for prompt_data in BENCHMARK_PROMPTS:
        prompt_id = prompt_data["id"]
        prompt = prompt_data["prompt"]
        max_tokens = prompt_data["max_tokens"]
        
        print(f"  [prompt] {prompt_id}...")
        
        # Format prompt with chat template
        messages = [{"role": "user", "content": prompt}]
        
        try:
            t0 = time.time()
            response = llm.create_chat_completion(
                messages=messages,
                max_tokens=max_tokens,
                temperature=0.7,
                top_p=0.9,
                stream=False,
            )
            elapsed = time.time() - t0
            
            text = response["choices"][0]["message"]["content"]
            actual_tokens = response["usage"]["completion_tokens"]
            prompt_tokens = response["usage"]["prompt_tokens"]
            
            tok_per_sec = actual_tokens / elapsed if elapsed > 0 else 0
            total_tokens += actual_tokens
            total_time += elapsed
            
            result = {
                "id": prompt_id,
                "category": prompt_data["category"],
                "description": prompt_data["description"],
                "prompt": prompt,
                "response": text,
                "response_length_chars": len(text),
                "completion_tokens": actual_tokens,
                "prompt_tokens": prompt_tokens,
                "elapsed_s": round(elapsed, 3),
                "tokens_per_second": round(tok_per_sec, 2),
                "time_to_first_token_s": None,  # Would need streaming
            }
            
            print(f"    → {actual_tokens} tokens in {elapsed:.2f}s = {tok_per_sec:.1f} tok/s")
            
        except Exception as e:
            print(f"    → FAILED: {e}")
            result = {
                "id": prompt_id,
                "error": str(e),
            }
        
        results["prompts"].append(result)
    
    # Calculate aggregate stats
    results["total_tokens"] = total_tokens
    results["total_generation_time_s"] = round(total_time, 2)
    results["avg_tokens_per_second"] = round(total_tokens / total_time, 2) if total_time > 0 else 0
    
    # Memory peak
    mem_peak = process.memory_info().rss / (1024 * 1024)
    results["memory_peak_mb"] = round(mem_peak, 1)
    
    # Cleanup
    print(f"  [cleanup] deleting model from memory...")
    del llm
    gc.collect()
    
    return results


def delete_model_file(model_path):
    """Delete the model file from disk to free space."""
    try:
        if os.path.exists(model_path):
            size = os.path.getsize(model_path) / (1024*1024)
            os.remove(model_path)
            print(f"  [cleanup] deleted {model_path} ({size:.0f} MB freed)")
    except Exception as e:
        print(f"  [cleanup] error deleting: {e}")


def run_full_benchmark():
    """Run benchmark on all models sequentially."""
    os.makedirs("/app/data/models", exist_ok=True)
    os.makedirs("/app/data/results", exist_ok=True)
    
    all_results = []
    
    for i, model_info in enumerate(MODELS):
        print(f"\n{'='*60}")
        print(f"BENCHMARKING MODEL {i+1}/{len(MODELS)}: {model_info['name']}")
        print(f"{'='*60}")
        
        model_dir = "/app/data/models"
        model_path = None
        
        try:
            # Step 1: Download
            model_path, dl_time = download_model(model_info, model_dir)
            
            # Step 2: Run benchmark
            results = run_model_benchmark(model_path, model_info)
            if results:
                results["download_time_s"] = round(dl_time, 2)
                all_results.append(results)
                
                # Save intermediate results
                with open(f"/app/data/results/benchmark_partial_{i:02d}_{model_info['name']}.json", "w") as f:
                    json.dump(results, f, indent=2, ensure_ascii=False)
                
                # Also upload to HF Space repo
                try:
                    from huggingface_hub import HfApi
                    api = HfApi()
                    api.upload_file(
                        path_or_fileobj=json.dumps(results, indent=2, ensure_ascii=False).encode(),
                        path_in_repo=f"results/benchmark_{model_info['name']}.json",
                        repo_id="TomatitoToho/zelin-benchmark",
                        repo_type="space",
                        token=os.environ.get("HF_TOKEN"),
                    )
                    print(f"  [upload] results saved to HF Space")
                except Exception as e:
                    print(f"  [upload] error: {e}")
            
        except Exception as e:
            print(f"  [error] model {model_info['name']} failed: {e}")
            import traceback
            traceback.print_exc()
            all_results.append({
                "model_info": model_info,
                "error": str(e),
            })
        
        finally:
            # Step 3: Delete model file
            if model_path:
                delete_model_file(model_path)
            
            # Force garbage collection
            gc.collect()
    
    # Save final combined results
    with open("/app/data/results/benchmark_final.json", "w") as f:
        json.dump({
            "total_models": len(all_results),
            "models_with_results": len([r for r in all_results if "prompts" in r]),
            "models_failed": len([r for r in all_results if "error" in r]),
            "results": all_results,
        }, f, indent=2, ensure_ascii=False)
    
    # Upload final
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        api.upload_file(
            path_or_fileobj=json.dumps({"results": all_results}, indent=2, ensure_ascii=False).encode(),
            path_in_repo="results/benchmark_final.json",
            repo_id="TomatitoToho/zelin-benchmark",
            repo_type="space",
            token=os.environ.get("HF_TOKEN"),
        )
        print(f"\n[upload] Final results uploaded to HF Space")
    except Exception as e:
        print(f"\n[upload] Error uploading final: {e}")
    
    print(f"\n{'='*60}")
    print(f"BENCHMARK COMPLETE")
    print(f"  Models tested: {len(all_results)}")
    print(f"  Models succeeded: {len([r for r in all_results if 'prompts' in r])}")
    print(f"  Models failed: {len([r for r in all_results if 'error' in r])}")
    print(f"{'='*60}")


if __name__ == "__main__":
    run_full_benchmark()