import json
import csv
import os
import subprocess
import time
import uuid
import urllib.request
import urllib.error

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
TASKS = os.path.join(ROOT, "benchmark", "locomo", "tasks.json")
OUT_CSV = os.path.join(ROOT, "benchmark", "locomo", "results.csv")
CLI_ARGS_BASE = ["run", "-p", "shcortex-cli", "--"]


def run_cmd(args):
    start = time.perf_counter()
    out = subprocess.check_output(["cargo"] + CLI_ARGS_BASE + args, text=True, cwd=ROOT)
    latency = time.perf_counter() - start
    return out.strip(), latency


def run_compress(text: str):
    return run_cmd(["compress", text])


def run_prompt(question: str):
    out, _ = run_cmd(["prompt", question])
    return out


def run_ollama(prompt: str):
    req = urllib.request.Request(
        "http://localhost:11434/api/generate",
        data=json.dumps({"model": "llama2", "prompt": prompt, "stream": False}).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    start = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            body = json.loads(resp.read().decode("utf-8"))
            answer = body.get("response", "")
    except Exception as e:
        answer = f"ERROR: {e}"
    latency = time.perf_counter() - start
    return answer, latency


def evaluate_answer(answer: str, expected: str) -> float:
    a = answer.lower()
    e = expected.lower()
    if not e:
        return 1.0
    tokens = e.split()
    matches = sum(1 for t in tokens if t in a)
    return matches / max(len(tokens), 1)


def main():
    with open(TASKS, "r", encoding="utf-8") as f:
        tasks = json.load(f)

    rows = []
    for task in tasks:
        conversation = "\n".join(task["conversation"])
        questions = task["questions"]
        expected = task.get("answers", [""] * len(questions))

        out1, lat1 = run_compress(conversation)

        for q_idx, q in enumerate(questions):
            opt_prompt = run_prompt(q)
            opt_resp, lat_opt = run_ollama(opt_prompt)

            exact = evaluate_answer(opt_resp, expected[q_idx] if q_idx < len(expected) else "")

            tokens_orig = len(conversation.split())
            tokens_opt = len(opt_prompt.split())
            compression = (tokens_orig - tokens_opt) / max(tokens_orig, 1)

            rows.append({
                "id": str(uuid.uuid4()),
                "task_id": task["id"],
                "lang": task["lang"],
                "question": q,
                "exact_match": exact,
                "compression": compression,
                "latency_shcortex_ms": lat1 * 1000,
                "latency_ollama_ms": lat_opt * 1000,
            })

    with open(OUT_CSV, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)

    print(f"Results saved to {OUT_CSV}")


if __name__ == "__main__":
    main()