import os
import sys
import json
import time
import argparse
from datetime import datetime

# Add root project path so we can import internal modules safely
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from GramAddict.core.llm_provider import query_telepathic_llm

BENCHMARKS_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "GramAddict/core/llm_benchmarks.json")
SCENARIOS_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "GramAddict/core/benchmark_scenarios.json")

def load_json(path):
    if os.path.exists(path):
        try:
            with open(path, "r") as f:
                return json.load(f)
        except Exception:
            return None
    return None

def save_json(path, data):
    with open(path, "w") as f:
        json.dump(data, f, indent=4)

def normalize_scores(db):
    if not db.get("models"):
        return db
    
    # 1. Find the highest raw score across all models
    max_raw = 0
    leader_model = None
    
    for name, data in db["models"].items():
        raw = data.get("raw_score", 0)
        if raw > max_raw:
            max_raw = raw
            leader_model = name
        elif raw == max_raw and max_raw > 0:
            # Tie-breaker: Latency
            current_lat = data.get("latency_ms", 99999)
            leader_lat = db["models"][leader_model].get("latency_ms", 99999)
            if current_lat < leader_lat:
                leader_model = name

    if max_raw == 0:
        return db

    # 2. Update relative performance
    for name, data in db["models"].items():
        raw = data.get("raw_score", 0)
        data["relative_performance_pct"] = round((raw / max_raw) * 100, 1)
        data["is_leader"] = (name == leader_model)
        
    return db

def benchmark_model(model_name: str, url: str, force: bool = False):
    db = load_json(BENCHMARKS_FILE) or {"models": {}}
    scenarios_data = load_json(SCENARIOS_FILE)
    if not scenarios_data:
        print("❌ Scenarios file missing!")
        return

    if not force and model_name in db.get("models", {}):
        pct = db["models"][model_name].get("relative_performance_pct", "N/A")
        print(f"Typical execution skip for {model_name} (Rel: {pct}%). Use --force.")
        return

    print(f"🚀 [Competitive Benchmarking] Model: {model_name}")
    
    total_raw = 0
    total_latency = 0
    results_detail = {}

    blank_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
    system_prompt = (
        "You identify which UI element to tap on an Android screen. "
        "Output ONLY valid JSON: {\"index\": number, \"reason\": \"brief reason\"}"
    )

    for scenario in scenarios_data["scenarios"]:
        print(f"--- Running: {scenario['name']} ---")
        
        user_prompt = (
            f"Which element should I tap to: {scenario['task']}\n\n"
            f"Elements:\n{json.dumps(scenario['nodes'], indent=1)}\n\n"
            "Rules:\n"
            "- Pick the SMALLEST, most specific button or icon\n"
            "- NEVER pick large containers\n"
            "Return: {\"index\": number, \"reason\": \"...\"}"
        )

        start_time = time.time()
        try:
            resp_str = query_telepathic_llm(model_name, url, system_prompt, user_prompt)
            latency = int((time.time() - start_time) * 1000)
            total_latency += latency
        except Exception as e:
            print(f"  ❌ API Request failed for scenario {scenario['id']}: {e}")
            continue

        raw_points = 0
        try:
            clean = resp_str.strip()
            if clean.startswith("```json"): clean = clean[7:]
            if clean.endswith("```"): clean = clean[:-3]
            data = json.loads(clean)
            
            # Points for structural adherence
            if "index" in data and "reason" in data:
                raw_points += 40
                
                # Points for correctness
                if data["index"] == scenario["target_index"]:
                    raw_points += 60
                    print(f"  ✅ Correct index ({data['index']}).")
                else:
                    print(f"  ❌ Wrong index ({data['index']}). Target was {scenario['target_index']}.")
            else:
                print("  ❌ JSON missing fields.")
        except Exception:
            print("  ❌ JSON Parsing failed.")

        results_detail[scenario["id"]] = raw_points
        total_raw += raw_points

    print(f"\n📊 Total Raw Score for {model_name}: {total_raw}")
    
    if model_name not in db["models"]:
        db["models"][model_name] = {}
        
    db["models"][model_name].update({
        "raw_score": total_raw,
        "latency_ms": total_latency // len(scenarios_data["scenarios"]),
        "last_tested": datetime.utcnow().isoformat() + "Z",
        "details": results_detail
    })
    
    # Recalculate relative scores across all models
    db = normalize_scores(db)
    save_json(BENCHMARKS_FILE, db)
    
    leader_name = [n for n, d in db["models"].items() if d.get("is_leader")][0]
    rel_pct = db["models"][model_name]["relative_performance_pct"]
    
    print(f"🏆 Current Leader: {leader_name}")
    print(f"✨ Relative Performance for {model_name}: {rel_pct}%")

if __name__ == "__main__":
    from GramAddict.core.config import Config
    
    parser = argparse.ArgumentParser(description="Competitive Benchmark for Singularity", add_help=False)
    parser.add_argument("--config", type=str, help="Bot config file")
    parser.add_argument("--model", type=str, help="Explicit model name")
    parser.add_argument("--url", type=str, help="Explicit endpoint URL")
    parser.add_argument("--force", action="store_true", help="Force re-testing")
    
    args, unknown = parser.parse_known_args()
    
    models_to_test = []
    
    if args.model and args.url:
        models_to_test.append((args.model, args.url))
    elif args.config:
        configs = Config(first_run=True, config=args.config)
        configs.parse_args()
        
        for attr, pref in [("ai_telepathic_model", "ai_telepathic_url"), ("ai_model", "ai_model_url"), ("ai_condenser_model", "ai_condenser_url")]:
            m = getattr(configs.args, attr, None)
            u = getattr(configs.args, pref, "https://openrouter.ai/api/v1/chat/completions")
            if m:
                models_to_test.append((m, u))
    else:
        print("❌ Syntax: --config test_config.yml or --model x --url y")
        sys.exit(1)
        
    for m, u in set(models_to_test):
        benchmark_model(m, u, args.force)
        time.sleep(1)