import os
import gc
import time
import json
import subprocess
import sys
import torch
from safetensors.torch import load_file
from eval_engine import MODEL_DIR, Qwen3FastRunner, decode_in_subprocess
from code_benchmark_problems import CODE_BENCHMARK_PROBLEMS, verify_code_solution

torch.set_num_threads(2)
CODE_TOK_CACHE_PATH = "/home/user/.cache/code_benchmark_tokens.pt"


def ensure_code_tokenized_cache():
    if os.path.exists(CODE_TOK_CACHE_PATH):
        return
    cmd = [
        sys.executable,
        "-c",
        (
            "import torch; from transformers import AutoTokenizer; "
            "from code_benchmark_problems import CODE_BENCHMARK_PROBLEMS; "
            f"tok = AutoTokenizer.from_pretrained('{MODEL_DIR}'); "
            "pids = [tok.encode(tok.apply_chat_template([{'role':'user','content':p['prompt']}], tokenize=False, add_generation_prompt=True, enable_thinking=False), add_special_tokens=False) for p in CODE_BENCHMARK_PROBLEMS]; "
            f"torch.save({{'prompt_ids_list': pids, 'eos_id': tok.eos_token_id, 'pad_id': tok.pad_token_id}}, '{CODE_TOK_CACHE_PATH}')"
        ),
    ]
    subprocess.run(cmd, check=True)


def evaluate_code_state_dict(sd, problems=CODE_BENCHMARK_PROBLEMS, max_new_tokens=90, batch_size=7, verbose=True):
    ensure_code_tokenized_cache()
    tok_data = torch.load(CODE_TOK_CACHE_PATH, weights_only=True)
    prompt_ids_list = tok_data["prompt_ids_list"]
    eos_id = tok_data["eos_id"]
    pad_id = tok_data["pad_id"]

    order = sorted(range(len(problems)), key=lambda i: len(prompt_ids_list[i]))
    runner = Qwen3FastRunner(sd)

    t0 = time.time()
    ordered_gen_ids = []
    num_batches = (len(problems) + batch_size - 1) // batch_size
    for b_idx, start in enumerate(range(0, len(problems), batch_size)):
        end = min(start + batch_size, len(problems))
        batch_indices = order[start:end]
        batch_prompts = [prompt_ids_list[i] for i in batch_indices]
        tb0 = time.time()
        gen_ids = runner.generate_batch(batch_prompts, eos_id, pad_id, max_new_tokens=max_new_tokens)
        ordered_gen_ids.extend(gen_ids)
        if verbose:
            avg_t = sum(len(g) for g in gen_ids) / len(gen_ids)
            print(f"  Batch {b_idx+1}/{num_batches} ({end-start} tasks, avg {avg_t:.1f} toks) in {time.time()-tb0:.1f}s", flush=True)
    elapsed = time.time() - t0

    all_gen_ids = [None] * len(problems)
    for orig_idx, g_ids in zip(order, ordered_gen_ids):
        all_gen_ids[orig_idx] = g_ids

    del runner
    gc.collect()

    responses = decode_in_subprocess(all_gen_ids)
    results = []
    passed = 0
    for prob, resp, g_ids in zip(problems, responses, all_gen_ids):
        ok, detail = verify_code_solution(prob, resp)
        if ok:
            passed += 1
        results.append({
            "id": prob["id"],
            "category": prob["category"],
            "prompt": prob["prompt"],
            "passed": ok,
            "detail": detail,
            "num_tokens": len(g_ids),
            "response": resp,
        })

    summary = {
        "total": len(problems),
        "passed": passed,
        "accuracy": round(100.0 * passed / len(problems), 2),
        "avg_tokens": round(sum(r["num_tokens"] for r in results) / len(results), 1),
        "elapsed_sec": round(elapsed, 2),
        "results": results,
    }
    return summary


if __name__ == "__main__":
    print("Evaluating Baseline Qwen3-0.6B on 28-task Executable Code Benchmark...", flush=True)
    sd = load_file(os.path.join(MODEL_DIR, "model.safetensors"), device="cpu")
    summary = evaluate_code_state_dict(sd, CODE_BENCHMARK_PROBLEMS, max_new_tokens=90, batch_size=7, verbose=True)
    with open("/home/user/code_baseline_results.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\n==========================================", flush=True)
    print(f"CODE BASELINE QWEN3-0.6B: {summary['passed']}/{summary['total']} ({summary['accuracy']}%) | Avg Toks: {summary['avg_tokens']} | {summary['elapsed_sec']}s", flush=True)
    print(f"==========================================", flush=True)
    for r in summary["results"]:
        status = "PASS" if r["passed"] else "FAIL"
        print(f"[{status:4s}] {r['id']:8s} ({r['category']:16s}) | Toks: {r['num_tokens']:3d} | {r['detail']}", flush=True)