import argparse
import json
import os
import subprocess
import sys
import time
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(__file__), "data/llm_benchmarks.json")
SCENARIOS_FILE = os.path.join(os.path.dirname(__file__), "data/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():
        if data.get("is_unsuitable"):
            continue

        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 get_installed_ollama_models():
    """
    Finds truly local Ollama models by parsing 'ollama list'.
    Strictly excludes remote/cloud endpoints or embedding-only models.
    """
    try:
        output = subprocess.check_output(["/usr/local/bin/ollama", "list"]).decode("utf-8")
        models = []
        for line in output.split("\n")[1:]:
            if line.strip():
                # Format: NAME, ID, SIZE, MODIFIED
                parts = line.split()
                if len(parts) >= 3:
                    name = parts[0]
                    size = parts[2]

                    # 1. Skip if size is '-' (remote/cloud model)
                    if size == "-":
                        continue

                    # 2. Skip ':cloud' tagged models explicitly
                    if ":cloud" in name:
                        continue

                    # 3. Filter out purely embedding models
                    if any(k in name.lower() for k in ["embed", "minilm", "rerank"]):
                        continue

                    models.append(name)
        return models
    except Exception as e:
        print(f"⚠️ Could not list Ollama models: {e}")
        return []


def benchmark_model(model_name: str, url: str, force: bool = False, iterations: int = 3):
    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")
        if not db["models"][model_name].get("is_unsuitable"):
            print(f"Typical execution skip for {model_name} (Rel: {pct}%). Use --force.")
            return

    print(f"\n🚀 [Competitive Benchmarking] Model: {model_name}")

    total_raw = 0
    total_latency = 0
    results_detail = {}
    passed_all = True

    system_prompt = (
        "You identify which UI element to tap based ONLY on a JSON array of parsed Android elements. "
        'Output ONLY valid JSON: {"index": number, "reason": "brief reason"}'
    )

    scenarios = scenarios_data["scenarios"]
    for scenario in 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\": \"...\"}"
        )

        scenario_latencies = []
        scenario_scores = []
        successes = 0

        for _ in range(iterations):
            start_time = time.time()
            try:
                resp_str = query_telepathic_llm(model_name, url, system_prompt, user_prompt)
                latency = int((time.time() - start_time) * 1000)
                scenario_latencies.append(latency)
            except Exception as e:
                print(f"  ❌ API Request failed for scenario {scenario['id']}: {e}")
                passed_all = False
                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
                        successes += 1
                    else:
                        print(f"  ❌ Wrong index ({data.get('index')}). Target was {scenario['target_index']}.")
                else:
                    print("  ❌ JSON missing fields.")
            except Exception:
                print("  ❌ JSON Parsing failed.")

            scenario_scores.append(raw_points)

        avg_scenario_score = int(sum(scenario_scores) / len(scenario_scores)) if scenario_scores else 0
        avg_scenario_latency = int(sum(scenario_latencies) / len(scenario_latencies)) if scenario_latencies else 0

        pass_rate = (successes / iterations) * 100
        if pass_rate < 100.0:
            passed_all = False

        print(
            f"  Result: {pass_rate:.0f}% Pass Rate | Avg Score: {avg_scenario_score}/100 | Avg Latency: {avg_scenario_latency}ms"
        )

        results_detail[scenario["id"]] = {
            "avg_score": avg_scenario_score,
            "pass_rate": pass_rate,
            "latency": avg_scenario_latency,
        }
        total_raw += avg_scenario_score
        total_latency += avg_scenario_latency

    avg_latency = total_latency // len(scenarios) if scenarios else 0
    print(
        f"\n📊 {model_name} Result: {'PASS' if passed_all else 'FAIL'} | Avg Score: {total_raw} | Latency: {avg_latency}ms"
    )

    if model_name not in db["models"]:
        db["models"][model_name] = {}

    db["models"][model_name].update(
        {
            "raw_score": total_raw,
            "telepathic_score": int((total_raw / (len(scenarios) * 100)) * 100) if scenarios else 0,
            "latency_ms": avg_latency,
            "last_tested": datetime.utcnow().isoformat() + "Z",
            "details": results_detail,
            "passed_all": passed_all,
            "is_unsuitable": not passed_all,
        }
    )

    # Recalculate relative scores across all models
    db = normalize_scores(db)
    save_json(BENCHMARKS_FILE, db)


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")
    parser.add_argument("--all-ollama", action="store_true", help="Automatically find and test all local Ollama models")
    parser.add_argument(
        "--iterations", type=int, default=3, help="Number of iterations per scenario to measure reliability"
    )

    args, unknown = parser.parse_known_args()

    models_to_test = []

    if args.all_ollama:
        ollama_models = get_installed_ollama_models()
        for m in ollama_models:
            models_to_test.append((m, "http://localhost:11434/api/generate"))
    elif 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, "http://localhost:11434/api/generate")
            if m:
                models_to_test.append((m, u))
    else:
        print("❌ Syntax: --all-ollama OR --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, args.iterations)
        time.sleep(1)
