import ollama
import time
import psutil
import os
import json

def get_vram_usage():
    try:
        # On macOS, we can check the 'ollama' process or use 'ollama ps'
        # For simplicity, we'll shell out to 'ollama ps' to get GPU residency
        result = os.popen('ollama ps').read()
        return result
    except:
        return "Unknown"

def benchmark_vision(model_name="moondream", image_path=None):
    print(f"\n--- Benchmarking Vision: {model_name} ---")
    if not image_path or not os.path.exists(image_path):
        print("No image provided for vision test.")
        return

    start_time = time.time()
    response = ollama.generate(
        model=model_name,
        prompt="Describe this image in detail.",
        images=[image_path],
        stream=False
    )
    duration = time.time() - start_time

    # Ollama returns tokens/sec in the response metadata
    eval_count = response.get('eval_count', 0)
    eval_duration = response.get('eval_duration', 1) / 1e9 # ns to s
    prompt_eval_count = response.get('prompt_eval_count', 0)
    prompt_eval_duration = response.get('prompt_eval_duration', 1) / 1e9

    print(f"Vision Response: {response['response'][:100]}...")
    print(f"Prompt Eval Rate: {prompt_eval_count / prompt_eval_duration:.2f} tokens/s")
    print(f"Generation Rate: {eval_count / eval_duration:.2f} tokens/s")
    print(f"VRAM Residency:\n{get_vram_usage()}")

def benchmark_context(model_name="qwen2-128k", context_size=131072):
    print(f"\n--- Benchmarking Context: {model_name} ({context_size} tokens) ---")
    
    # Generate a long dummy context
    # Each 'word' is approx 1.3 tokens
    dummy_text = "This is a context verification test. " * (context_size // 10)
    
    start_time = time.time()
    response = ollama.generate(
        model=model_name,
        prompt=f"{dummy_text}\n\nTask: What was the first sentence of this text?",
        stream=False
    )
    duration = time.time() - start_time

    eval_count = response.get('eval_count', 0)
    eval_duration = response.get('eval_duration', 1) / 1e9
    prompt_eval_count = response.get('prompt_eval_count', 0)
    prompt_eval_duration = response.get('prompt_eval_duration', 1) / 1e9

    print(f"Context Response: {response['response']}")
    print(f"Prompt Eval Rate (Context Parsing): {prompt_eval_count / prompt_eval_duration:.2f} tokens/s")
    print(f"Generation Rate: {eval_count / eval_duration:.2f} tokens/s")
    print(f"VRAM Residency:\n{get_vram_usage()}")

if __name__ == "__main__":
    # Ensure we have a sample image
    # For CI/Automation, we'll try to find any image in the repo or use a placeholder
    sample_image = "/Users/MWsatwareAG/internal/mw-macbook-pro/ml-test/sample.png"
    # Create a small dummy image if it doesn't exist for the script to run
    if not os.path.exists(sample_image):
        import matplotlib.pyplot as plt
        import numpy as np
        plt.imshow(np.random.rand(100,100,3))
        plt.savefig(sample_image)
        print(f"Created dummy image at {sample_image}")

    benchmark_vision(model_name="moondream", image_path=sample_image)
    # 128k context test
    try:
        benchmark_context(model_name="qwen2-128k-gpu")
    except Exception as e:
        print(f"Context test failed (likely VRAM/OOM): {e}")