import json
import multiprocessing
import os
import re
import sys
import time

os.environ["CUDA_HOME"] = "/usr/local/cuda-13.0"
os.environ["PATH"] = f"/usr/local/cuda-13.0/bin:{os.environ.get('PATH', '')}"

MODEL_DIR = "/home/user/models/DeepSeek-V4-Flash"
sys.path.insert(0, os.path.join(MODEL_DIR, "encoding"))
from encoding_dsv4 import encode_messages

from datasets import load_dataset
from vllm import LLM, SamplingParams

OUTPUT_FILE = sys.argv[1] if len(sys.argv) > 1 else "/home/user/humaneval_results_deepseek_v4.jsonl"
THINKING_MODE = "thinking"


def check_correctness(full_code: str, timeout: int = 15):

    def run_test(result_dict):
        try:
            exec_globals = {}
            exec(full_code, exec_globals)
            result_dict["passed"] = True
        except Exception:
            result_dict["passed"] = False

    manager = multiprocessing.Manager()
    result_dict = manager.dict()
    proc = multiprocessing.Process(target=run_test, args=(result_dict,))
    proc.start()
    proc.join(timeout)
    if proc.is_alive():
        proc.kill()
        manager.shutdown()
        return False
    passed = result_dict.get("passed", False)
    manager.shutdown()
    return passed


def extract_code_block(text: str) -> str | None:
    blocks = re.findall(r"```(?:python)?\s*\n(.*?)```", text, re.DOTALL)
    if blocks:
        for b in blocks:
            b = b.strip()
            if "def " in b or "import " in b or b.startswith("from "):
                return b
        return blocks[-1].strip()
    return None


def main():
    print("Loading HumanEval dataset...")
    ds = load_dataset("openai_humaneval", split="test")
    print(f"Loaded {len(ds)} problems")

    print("Loading model with vLLM...")
    llm = LLM(
        model=MODEL_DIR,
        tensor_parallel_size=1,
        dtype="auto",
        kv_cache_dtype="fp8",
        max_model_len=32768,
        trust_remote_code=True,
    )

    sampling_params = SamplingParams(
        temperature=0.0,
        top_p=0.95,
        max_tokens=4096,
        stop=["<｜end▁of▁sentence｜>"],
    )

    formatted_prompts = []
    metadata = []
    for example in ds:
        prompt = example["prompt"]
        messages = [
            {
                "role": "user",
                "content": f"Write a Python function to solve the following problem:\n\n{prompt}",
            },
        ]
        formatted = encode_messages(messages, thinking_mode=THINKING_MODE)
        formatted_prompts.append(formatted)
        metadata.append({
            "task_id": example["task_id"],
            "entry_point": example["entry_point"],
            "test": example["test"],
            "prompt": example["prompt"],
        })

    print(f"Sample prompt: {formatted_prompts[0][:200]}...")
    print(f"Generating completions for {len(formatted_prompts)} problems...")
    start = time.time()

    outputs = llm.generate(formatted_prompts, sampling_params)

    elapsed = time.time() - start
    print(f"Generation completed in {elapsed:.2f}s ({elapsed / len(outputs):.2f}s per sample)")

    results = []
    for out, meta in zip(outputs, metadata):
        raw = out.outputs[0].text.strip()
        if THINKING_MODE == "thinking":
            end_idx = raw.find("</think>")
            if end_idx != -1:
                code_text = raw[end_idx + len("</think>"):].strip()
            else:
                code_text = raw
        else:
            code_text = raw
        code = extract_code_block(code_text)
        if not code:
            code = code_text
        code = re.sub(r"^```(?:python)?\s*", "", code)
        code = re.sub(r"\s*```$", "", code)
        code = code.strip()
        results.append({
            "task_id": meta["task_id"],
            "entry_point": meta["entry_point"],
            "prompt": meta["prompt"],
            "test": meta["test"],
            "generation": code,
            "raw_output": raw,
        })

    with open(OUTPUT_FILE, "w") as f:
        for r in results:
            f.write(json.dumps(r) + "\n")
    print(f"Results saved to {OUTPUT_FILE}")

    passed = 0
    for r in results:
        gen = r["generation"]
        if f"def {r['entry_point']}" in gen:
            full_code = f"{gen}\n{r['test']}\ncheck({r['entry_point']})"
        else:
            full_code = f"{r['prompt']}\n{gen}\n{r['test']}\ncheck({r['entry_point']})"
        if check_correctness(full_code):
            passed += 1

    total = len(results)
    print(f"\nPass@1: {passed}/{total} = {passed / total * 100:.2f}%")


if __name__ == "__main__":
    main()