from __future__ import annotations

import argparse
import fcntl
import json
import os
import random
import re
import sys
import types
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path("/225040511/project/Biomni-ReAct")
LAB_BENCH_ROOT = Path("/225040511/project/LAB-Bench")
DEFAULT_OUTPUT_ROOT = PROJECT_ROOT / "LAB-bench"
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)
LETTER_RE = re.compile(r"\b([A-Z])\b", re.IGNORECASE)

sys.path.insert(0, str(PROJECT_ROOT))
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

from biomni_react.agent import BiomniReActAgent  # noqa: E402
from biomni_react.config import AgentConfig  # noqa: E402
from biomni_react.schema import TaskSpec  # 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(
    *,
    eval_name: str,
    split: str,
    dev_size: int,
    test_size: int,
    seed: int,
    shard_index: int,
    shard_count: int,
    debug: bool,
) -> list[tuple[str, Any]]:
    evaluator = load_eval(eval_name)
    instances = list(evaluator.eval_set.instances)
    rng = random.Random(f"{seed}:{eval_name}:question-set")
    rng.shuffle(instances)
    if debug:
        selected = instances[: min(3, len(instances))]
    elif split == "dev":
        selected = instances[: min(dev_size, len(instances))]
    elif split == "test":
        start = min(dev_size, len(instances))
        selected = instances[start : min(start + test_size, len(instances))]
    else:
        selected = instances
    if shard_count > 1:
        total = len(selected)
        chunk_size = (total + shard_count - 1) // shard_count
        selected = selected[min(total, shard_index * chunk_size) : min(total, (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_questions: set[str] = set()
    completed_task_ids: 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("task_id") or "").strip()
        if task_id:
            completed_task_ids.add(task_id)
        question = str(record.get("question") or "").strip()
        if question:
            completed_questions.add(question)
    return completed_task_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_objective(input_obj: Any, eval_name: str) -> str:
    choices = "\n".join(input_obj.choices)
    return f"""
Answer this multiple-choice LAB-Bench biology question from {eval_name}.

Question:
{input_obj.question}

Options:
{choices}

Return the single correct letter. You must write answer.txt containing exactly:
[ANSWER]X[/ANSWER]
where X is one answer letter.

Finish with <solution>[ANSWER]X[/ANSWER]</solution>.
""".strip()


def make_agent_config(args: argparse.Namespace) -> AgentConfig:
    api_key = (
        args.api_key
        or os.getenv("BIOMNI_REACT_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, BIOMNI_REACT_API_KEY, BIOMNI_CUSTOM_API_KEY, or OPENAI_API_KEY.")
    return AgentConfig(
        model=args.model,
        base_url=args.base_url,
        api_key=api_key,
        max_iterations=args.max_iterations,
        retrieval_top_k=args.top_k,
        command_timeout_s=args.command_timeout,
    )


def run_one(
    *,
    agent: BiomniReActAgent,
    eval_name: str,
    split: str,
    subset: str,
    instance: Any,
    output_root: Path,
    result_file: Path,
    reasoning_log: Path,
    method: str,
    model: str,
) -> dict[str, Any]:
    input_obj, target_output, _unsure = instance.get_input_output()
    case_dir = output_root / "case_workspaces" / f"{eval_name.lower()}_{instance.id}"
    answer_path = case_dir / "answer.txt"
    task = TaskSpec(
        name=f"LAB-Bench {eval_name} {instance.id}",
        objective=build_objective(input_obj, eval_name),
        workspace=case_dir,
        expected_outputs=[answer_path],
        constraints=[
            "Do not use external network resources.",
            "Do not inspect answer keys or previous result files.",
            "Write exactly one answer letter wrapped in [ANSWER] and [/ANSWER].",
        ],
        metadata={
            "eval": eval_name,
            "split": split,
            "subset": subset,
            "question_id": str(instance.id),
        },
    )
    try:
        result = agent.run(task)
        raw_output = result.final_answer
        if answer_path.exists():
            raw_output = answer_path.read_text(encoding="utf-8", errors="replace") + "\n" + raw_output
        error = result.error or ""
    except Exception as exc:
        raw_output = ""
        error = repr(exc)
    answer = parse_answer(raw_output, len(input_obj.choices))
    record = {
        "task_id": str(instance.id),
        "subset": subset,
        "question": str(input_obj.question),
        "answer": str(target_output),
        "agent_answer": answer,
        "method": method,
        "model": model,
        "error": error,
    }
    append_jsonl_locked(result_file, record)
    append_text_locked(
        reasoning_log,
        "\n".join(
            [
                "=" * 80,
                f"eval: {eval_name}",
                f"split: {split}",
                f"id: {instance.id}",
                f"subset: {subset}",
                f"answer: {target_output}",
                f"agent_answer: {answer}",
                f"workspace: {case_dir}",
                f"error: {error}",
                "",
                "[raw_output]",
                raw_output,
                "",
            ]
        ),
    )
    return record


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run LAB-Bench with Biomni-ReAct.")
    parser.add_argument("--eval", choices=[member.value for member in labbench.Eval], default="DbQA")
    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("BIOMNI_REACT_MODEL", os.getenv("DEEPSEEK_MODEL_NAME", "deepseek-chat")))
    parser.add_argument("--base-url", default=os.getenv("BIOMNI_REACT_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("BIOMNI_REACT_MAX_ITERATIONS", "8")))
    parser.add_argument("--top-k", type=int, default=int(os.getenv("BIOMNI_REACT_TOP_K", "8")))
    parser.add_argument("--command-timeout", type=int, default=int(os.getenv("BIOMNI_REACT_TIMEOUT", "120")))
    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()
    result_file = args.result_file or args.output_root / f"{eval_lower}_results.jsonl"
    reasoning_log = args.reasoning_log or args.output_root / f"{eval_lower}_reasoning.log"
    selected = select_instances(
        eval_name=args.eval,
        split=args.split,
        dev_size=args.dev_size,
        test_size=args.test_size,
        seed=args.seed,
        shard_index=args.shard_index,
        shard_count=args.shard_count,
        debug=args.debug,
    )
    if args.skip_existing_results:
        completed_task_ids, completed_questions = load_completed_results(result_file)
        selected = [
            (subset, instance)
            for subset, instance in selected
            if str(instance.id) not in completed_task_ids
            and str(instance.get_input_output()[0].question).strip() not in completed_questions
        ]
    config = make_agent_config(args)
    agent = BiomniReActAgent(config=config)
    for subset, instance in selected:
        record = run_one(
            agent=agent,
            eval_name=args.eval,
            split=args.split,
            subset=subset,
            instance=instance,
            output_root=args.output_root,
            result_file=result_file,
            reasoning_log=reasoning_log,
            method="Biomni-ReAct",
            model=config.model,
        )
        print(json.dumps(record, ensure_ascii=False), flush=True)
    return 0


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