#!/usr/bin/env python3
from __future__ import annotations

import argparse
import contextlib
import fcntl
import io
import json
import os
import random
import re
import subprocess
import sys
import tempfile
import textwrap
import types
from pathlib import Path
from typing import Any


PROJECT_ROOT = Path("/225040511/project/React+code_labbench")
LAB_BENCH_ROOT = Path("/225040511/project/LAB-Bench")
DEFAULT_OUTPUT_ROOT = PROJECT_ROOT / "results"
DEFAULT_DEV_SIZE = 45
DEFAULT_TEST_SIZE = 315
DEFAULT_SEED = 20260514

ANSWER_RE = re.compile(r"\[ANSWER\]\s*([A-Z])\s*\[/ANSWER\]", re.IGNORECASE)
SOLUTION_RE = re.compile(r"<solution>\s*(.*?)\s*</solution>", re.IGNORECASE | re.DOTALL)
PY_BLOCK_RE = re.compile(r"```(?:python)?\s*(.*?)```", re.IGNORECASE | re.DOTALL)
EXEC_RE = re.compile(r"<execute(?:\s+type=[\"']python[\"'])?\s*>(.*?)</execute>", re.IGNORECASE | re.DOTALL)
LETTER_RE = re.compile(r"\b([A-Z])\b", re.IGNORECASE)

sys.path.insert(0, str(LAB_BENCH_ROOT))


def install_labbench_import_stubs() -> None:
    if "vertexai" not in sys.modules:
        vertexai = types.ModuleType("vertexai")
        vertexai.init = lambda *_args, **_kwargs: None
        sys.modules["vertexai"] = vertexai
    if "google.auth" not in sys.modules:
        google = sys.modules.setdefault("google", types.ModuleType("google"))
        auth = types.ModuleType("google.auth")
        auth.default = lambda *_args, **_kwargs: (types.SimpleNamespace(refresh=lambda *_a, **_k: None, token=""), None)
        transport = types.ModuleType("google.auth.transport")
        requests = types.ModuleType("google.auth.transport.requests")
        requests.Request = lambda *_args, **_kwargs: None
        transport.requests = requests
        auth.transport = transport
        google.auth = auth
        sys.modules["google.auth"] = auth
        sys.modules["google.auth.transport"] = transport
        sys.modules["google.auth.transport.requests"] = requests
    if "chembench" not in sys.modules:
        chembench = types.ModuleType("chembench")
        sys.modules["chembench"] = chembench
        constant = types.ModuleType("chembench.constant")
        constant.COT_PROMPT = "Think step by step."
        constant.MCQ_REGEX_TEMPLATE_1 = r"\[ANSWER\]\s*([A-Z])\s*\[/ANSWER\]"
        sys.modules["chembench.constant"] = constant
        prompter = types.ModuleType("chembench.prompter")
        prompter.prepare_mcq_answer = lambda text, *_args, **_kwargs: text
        sys.modules["chembench.prompter"] = prompter
        utils = types.ModuleType("chembench.utils")
        utils.create_multiple_choice_regex = lambda letters: r"\b(" + "|".join(letters) + r")\b"
        utils.post_process_prompts = lambda text: text
        utils.run_regex = lambda _regex, text, return_first=True: None
        sys.modules["chembench.utils"] = utils


install_labbench_import_stubs()
import labbench  # noqa: E402


def load_dotenv_files(paths: list[Path]) -> None:
    for path in paths:
        if not path.exists():
            continue
        for raw_line in path.read_text(encoding="utf-8", errors="replace").splitlines():
            line = raw_line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, value = line.split("=", 1)
            os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))


def load_eval(eval_name: str) -> labbench.Evaluator:
    return labbench.Evaluator(labbench.Eval(eval_name), debug=False, open_answer=False, use_hf=False)


def select_instances(args: argparse.Namespace) -> list[tuple[str, Any]]:
    evaluator = load_eval(args.eval)
    instances = list(evaluator.eval_set.instances)
    rng = random.Random(f"{args.seed}:{args.eval}:question-set")
    rng.shuffle(instances)
    if args.debug:
        selected = instances[: min(3, len(instances))]
    elif args.split == "dev":
        selected = instances[: min(args.dev_size, len(instances))]
    elif args.split == "test":
        start = min(args.dev_size, len(instances))
        selected = instances[start : min(start + args.test_size, len(instances))]
    else:
        selected = instances
    if args.shard_count > 1:
        total = len(selected)
        chunk_size = (total + args.shard_count - 1) // args.shard_count
        selected = selected[
            min(total, args.shard_index * chunk_size) : min(total, (args.shard_index + 1) * chunk_size)
        ]
    return selected


def load_completed_results(path: Path) -> tuple[set[str], set[str]]:
    if not path.exists():
        return set(), set()
    completed_ids: set[str] = set()
    completed_questions: set[str] = set()
    for raw_line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        if not raw_line.strip():
            continue
        try:
            record = json.loads(raw_line)
        except json.JSONDecodeError:
            continue
        task_id = str(record.get("id") or record.get("task_id") or "").strip()
        if task_id:
            completed_ids.add(task_id)
        question = str(record.get("question") or "").strip()
        if question:
            completed_questions.add(question)
    return completed_ids, completed_questions


def append_text_locked(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
        handle.write(text)
        handle.flush()
        os.fsync(handle.fileno())
        fcntl.flock(handle.fileno(), fcntl.LOCK_UN)


def append_jsonl_locked(path: Path, payload: dict[str, Any]) -> None:
    append_text_locked(path, json.dumps(payload, ensure_ascii=False, default=str) + "\n")


def parse_answer(text: str, n_choices: int) -> str:
    valid = set("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[:n_choices])
    if match := ANSWER_RE.search(text or ""):
        letter = match.group(1).upper()
        if letter in valid:
            return letter
    if match := SOLUTION_RE.search(text or ""):
        return parse_answer(match.group(1), n_choices)
    for match in LETTER_RE.finditer(text or ""):
        letter = match.group(1).upper()
        if letter in valid:
            return letter
    return ""


def build_prompt(input_obj: Any, eval_name: str) -> str:
    choices = "\n".join(input_obj.choices)
    return f"""
You are a biology reasoning agent with a Python scratchpad.
Answer this multiple-choice LAB-Bench question from {eval_name}.

You may reason and write Python code for sequence analysis. If you need code,
return it in either <execute>...</execute> or a Python fenced block.
When finished, return exactly: <solution>[ANSWER]X[/ANSWER]</solution>

Question:
{input_obj.question}

Choices:
{choices}
""".strip()


def execute_python(code: str, timeout: int) -> str:
    with tempfile.TemporaryDirectory(prefix="react_code_labbench_") as tmp:
        script = Path(tmp) / "scratch.py"
        script.write_text(code, encoding="utf-8")
        try:
            completed = subprocess.run(
                [sys.executable, str(script)],
                cwd=tmp,
                text=True,
                capture_output=True,
                timeout=timeout,
            )
        except subprocess.TimeoutExpired:
            return "Execution timed out."
        output = completed.stdout
        if completed.stderr:
            output += "\n[stderr]\n" + completed.stderr
        return output[-12000:]


def resolve_client(args: argparse.Namespace):
    try:
        from openai import OpenAI
    except ImportError as exc:
        raise SystemExit("openai package is required for React+code runner.") from exc
    api_key = (
        args.api_key
        or os.getenv("REACT_CODE_API_KEY")
        or os.getenv("DEEPSEEK_API_KEY")
        or os.getenv("BIOMNI_CUSTOM_API_KEY")
        or os.getenv("OPENAI_API_KEY")
    )
    if not api_key:
        raise SystemExit("Missing API key. Set DEEPSEEK_API_KEY, REACT_CODE_API_KEY, BIOMNI_CUSTOM_API_KEY, or OPENAI_API_KEY.")
    return OpenAI(api_key=api_key, base_url=args.base_url), args.model


def react_code_answer(client: Any, model: str, prompt: str, max_iters: int, timeout: int) -> tuple[str, str, dict[str, int]]:
    messages = [
        {
            "role": "system",
            "content": "Use concise reasoning. Use Python when helpful. Return a final <solution>[ANSWER]X[/ANSWER]</solution>.",
        },
        {"role": "user", "content": prompt},
    ]
    transcript = []
    usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "llm_call_count": 0}
    for _ in range(max_iters):
        response = client.chat.completions.create(model=model, temperature=0, messages=messages)
        usage["llm_call_count"] += 1
        if getattr(response, "usage", None):
            usage["prompt_tokens"] += int(getattr(response.usage, "prompt_tokens", 0) or 0)
            usage["completion_tokens"] += int(getattr(response.usage, "completion_tokens", 0) or 0)
            usage["total_tokens"] += int(getattr(response.usage, "total_tokens", 0) or 0)
        content = response.choices[0].message.content or ""
        transcript.append(content)
        if ANSWER_RE.search(content) or SOLUTION_RE.search(content):
            return content, "\n\n".join(transcript), usage
        code = None
        if match := EXEC_RE.search(content):
            code = match.group(1).strip()
        elif match := PY_BLOCK_RE.search(content):
            code = match.group(1).strip()
        if not code:
            messages.append({"role": "assistant", "content": content})
            messages.append({"role": "user", "content": "Finish with <solution>[ANSWER]X[/ANSWER]</solution>."})
            continue
        observation = execute_python(code, timeout=timeout)
        transcript.append("[observation]\n" + observation)
        messages.append({"role": "assistant", "content": content})
        messages.append({"role": "user", "content": "Python observation:\n" + observation})
    return "", "\n\n".join(transcript), usage


def run_one(args: argparse.Namespace, client: Any, model: str, subset: str, instance: Any) -> dict[str, Any]:
    input_obj, target_output, _unsure = instance.get_input_output()
    final, transcript, usage = react_code_answer(
        client,
        model,
        build_prompt(input_obj, args.eval),
        max_iters=args.max_iterations,
        timeout=args.command_timeout,
    )
    answer = parse_answer(final or transcript, len(input_obj.choices))
    record = {
        "id": str(instance.id),
        "task_id": str(instance.id),
        "subset": subset,
        "question": str(input_obj.question),
        "answer": str(target_output),
        "agent_answer": answer,
        "method": "React+code",
        "model": model,
        **usage,
    }
    append_jsonl_locked(args.result_file, record)
    append_text_locked(
        args.reasoning_log,
        "\n".join(
            [
                "=" * 80,
                f"eval: {args.eval}",
                f"split: {args.split}",
                f"id: {instance.id}",
                f"answer: {target_output}",
                f"agent_answer: {answer}",
                "",
                "[transcript]",
                transcript,
                "",
            ]
        ),
    )
    return record


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run LAB-Bench with a lightweight React+code baseline.")
    parser.add_argument("--eval", choices=[member.value for member in labbench.Eval], default="CloningScenarios")
    parser.add_argument("--split", choices=["dev", "test", "all"], default="test")
    parser.add_argument("--dev-size", type=int, default=DEFAULT_DEV_SIZE)
    parser.add_argument("--test-size", type=int, default=DEFAULT_TEST_SIZE)
    parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--shard-count", type=int, default=1)
    parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT)
    parser.add_argument("--result-file", type=Path, default=None)
    parser.add_argument("--reasoning-log", type=Path, default=None)
    parser.add_argument("--skip-existing-results", action="store_true")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--model", default=os.getenv("REACT_CODE_MODEL", os.getenv("DEEPSEEK_MODEL_NAME", "deepseek-chat")))
    parser.add_argument("--base-url", default=os.getenv("REACT_CODE_BASE_URL", os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")))
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--max-iterations", type=int, default=int(os.getenv("REACT_CODE_MAX_ITERATIONS", "6")))
    parser.add_argument("--command-timeout", type=int, default=int(os.getenv("REACT_CODE_TIMEOUT", "60")))
    parser.add_argument("--env-file", action="append", type=Path, default=[])
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.shard_count < 1:
        raise SystemExit("--shard-count must be at least 1.")
    if args.shard_index < 0 or args.shard_index >= args.shard_count:
        raise SystemExit("--shard-index must satisfy 0 <= shard-index < shard-count.")
    load_dotenv_files([PROJECT_ROOT / ".env", LAB_BENCH_ROOT / ".env", Path("/225040511/project/.env"), *args.env_file])
    args.output_root.mkdir(parents=True, exist_ok=True)
    eval_lower = args.eval.lower()
    args.result_file = args.result_file or args.output_root / f"{eval_lower}_results.jsonl"
    args.reasoning_log = args.reasoning_log or args.output_root / f"{eval_lower}_reasoning.log"
    selected = select_instances(args)
    if args.skip_existing_results:
        completed_ids, completed_questions = load_completed_results(args.result_file)
        selected = [
            (subset, instance)
            for subset, instance in selected
            if str(instance.id) not in completed_ids
            and str(instance.get_input_output()[0].question).strip() not in completed_questions
        ]
    client, model = resolve_client(args)
    for subset, instance in selected:
        print(json.dumps(run_one(args, client, model, subset, instance), ensure_ascii=False), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())