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_llm, 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")

# Minimum iterations for statistical significance
MIN_ITERATIONS = 5


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):
    """Normalize relative performance by AVERAGE score per scenario, not raw totals."""
    if not db.get("models"):
        return db

    max_avg = 0
    leader_model = None

    for name, data in db["models"].items():
        if data.get("is_unsuitable"):
            continue

        scenario_count = data.get("scenario_count", 1)
        avg = data.get("raw_score", 0) / max(scenario_count, 1)
        data["avg_score_per_scenario"] = round(avg, 1)

        if avg > max_avg:
            max_avg = avg
            leader_model = name
        elif avg == max_avg and max_avg > 0:
            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_avg == 0:
        return db

    for name, data in db["models"].items():
        scenario_count = data.get("scenario_count", 1)
        avg = data.get("raw_score", 0) / max(scenario_count, 1)
        data["relative_performance_pct"] = round((avg / max_avg) * 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():
                parts = line.split()
                if len(parts) >= 3:
                    name = parts[0]
                    size = parts[2]

                    if size == "-":
                        continue
                    if ":cloud" in name:
                        continue
                    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 _run_telepathic_scenario(scenario, model_name, url, iterations):
    """Run a telepathic (JSON element selection) scenario."""
    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"}'
    )

    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": "..."}'
    )

    latencies = []
    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)
            latencies.append(latency)
        except Exception as e:
            print(f"  ❌ API Request failed: {e}")
            scores.append(0)
            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)

            if "index" in data and "reason" in data:
                raw_points += 40
                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.")

        scores.append(raw_points)

    return scores, latencies, successes


def _run_brain_scenario(scenario, model_name, url, iterations):
    """Run a brain action extraction scenario (format_json=False)."""
    system_prompt = (
        f"You are an autonomous Instagram agent. Your goal is: '{scenario['task']}'.\n"
        f"You are currently on screen: {scenario['screen_type']}.\n"
        f"Available actions: {scenario['available_actions']}\n"
        "INSTRUCTIONS: Reply with ONLY the action string. Nothing else."
    )

    user_prompt = "Choose the next best action."

    latencies = []
    scores = []
    successes = 0

    for _ in range(iterations):
        start_time = time.time()
        try:
            # CRITICAL: Use format_json=False — this is the Brain code path
            ans = query_llm(
                url=url,
                model=model_name,
                prompt=user_prompt,
                system=system_prompt,
                format_json=False,
                timeout=30,
                temperature=0.0,
                max_tokens=50,
            )
            latency = int((time.time() - start_time) * 1000)
            latencies.append(latency)
        except Exception as e:
            print(f"  ❌ API Request failed: {e}")
            scores.append(0)
            continue

        raw_points = 0
        if ans and "response" in ans:
            response = ans["response"].strip().lower()

            # Points for structural adherence (returned a clean string)
            if response and response in [a.lower() for a in scenario["available_actions"]]:
                raw_points += 40

                # Points for correctness
                if scenario.get("accept_any_valid"):
                    # Any valid action from the list is acceptable
                    raw_points += 60
                    successes += 1
                elif response == scenario["target_action"].lower():
                    raw_points += 60
                    successes += 1
                else:
                    print(f"  ⚠️  Valid but suboptimal: '{response}' (target: '{scenario['target_action']}')")
                    raw_points += 20  # Partial credit for valid but wrong action
            else:
                print(f"  ❌ Invalid response: '{response}' not in available actions")
        else:
            print("  ❌ Empty or null response from LLM")

        scores.append(raw_points)

    return scores, latencies, successes


def benchmark_model(model_name: str, url: str, force: bool = False, iterations: int = MIN_ITERATIONS):
    iterations = max(iterations, MIN_ITERATIONS)  # Enforce minimum

    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} ({iterations} iterations)")

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

    scenarios = scenarios_data["scenarios"]
    for scenario in scenarios:
        scenario_type = scenario.get("type", "telepathic")
        print(f"--- [{scenario_type.upper()}] {scenario['name']} ---")

        if scenario_type == "telepathic":
            scores, latencies, successes = _run_telepathic_scenario(scenario, model_name, url, iterations)
        elif scenario_type == "brain_action":
            scores, latencies, successes = _run_brain_scenario(scenario, model_name, url, iterations)
        else:
            print(f"  ⚠️ Unknown scenario type: {scenario_type}")
            continue

        avg_score = int(sum(scores) / len(scores)) if scores else 0
        avg_latency = int(sum(latencies) / len(latencies)) if latencies else 0
        pass_rate = (successes / iterations) * 100

        if pass_rate < 100.0:
            passed_all = False

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

        # Consistent format: always an object
        results_detail[scenario["id"]] = {
            "avg_score": avg_score,
            "pass_rate": pass_rate,
            "latency": avg_latency,
        }
        total_raw += avg_score
        total_latency += avg_latency

    avg_latency = total_latency // len(scenarios) if scenarios else 0
    print(f"\n📊 {model_name}: {'PASS' if passed_all else 'FAIL'} | Total: {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,
            "scenario_count": len(scenarios),
            "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,
            "iterations": iterations,
        }
    )

    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=MIN_ITERATIONS, help=f"Iterations per scenario (min: {MIN_ITERATIONS})"
    )

    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)
