#!/usr/bin/env python3
"""Convert PlantGeneBench 100 Markdown tables into full structured JSONL/gold."""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any

LAYER_BY_CODE = {
    "TU": "task_understanding",
    "EP": "execution_planning",
    "EC": "evidence_constraint",
    "BI": "biological_interpretation",
    "BC": "boundary_control",
}

PROFILE_BY_LAYER = {
    "task_understanding": ["intent", "entities", "target_stage", "boundary"],
    "execution_planning": ["target_stage", "execution_plan", "parameters", "artifacts", "boundary"],
    "evidence_constraint": ["target_stage", "evidence", "boundary", "hallucination"],
    "biological_interpretation": ["evidence", "biology", "boundary", "hallucination"],
    "boundary_control": ["boundary", "hallucination", "target_stage"],
}

SPECIES_ALIASES = {
    "rice": "rice",
    "Arabidopsis": "Arabidopsis",
    "maize": "maize",
    "soybean": "soybean",
    "tomato": "tomato",
    "Fagopyrum_tataricum": "Fagopyrum_tataricum",
}

ACTION_PATTERNS = [
    "不要安排",
    "不要运行",
    "不要重新跑",
    "不要加入",
    "不要调用",
    "不要转到",
    "不要启动",
    "不要强制",
    "不要直接",
    "不要只运行",
    "不要创建",
]

STAGE_TERMS = {
    "identify": ["identify"],
    "phylogeny": ["phylogeny"],
    "motif": ["motif"],
    "domain": ["domain"],
    "structure": ["structure"],
    "promoter": ["promoter"],
    "synteny": ["synteny"],
    "interspecies_synteny": ["interspecies_synteny"],
    "karyotype": ["karyotype"],
    "combo": ["combo"],
    "ppi": ["PPI", "ppi"],
    "expression": ["expression"],
    "pep": ["pep"],
    "rnaseq_qc": ["rnaseq_qc"],
    "rnaseq_align": ["rnaseq_align"],
    "rnaseq_quant": ["rnaseq_quant"],
    "rnaseq_sample_qc": ["rnaseq_sample_qc"],
    "rnaseq_de": ["rnaseq_de"],
    "rnaseq_enrich": ["rnaseq_enrich", "rnaseq_go_enrich", "rnaseq_kegg_enrich"],
    "rnaseq_wgcna": ["rnaseq_wgcna"],
    "rnaseq_fgsea": ["rnaseq_fgsea"],
    "rnaseq_candidate_family": ["rnaseq_candidate_family"],
    "reference_classification": ["reference_classification"],
    "subfamily_confirm": ["subfamily_confirm"],
    "ortholog": ["ortholog"],
    "kaks": ["kaks"],
    "summary": ["summary"],
    "annotation": ["annotation"],
    "tree_plot": ["tree_plot"],
}

ANALYSIS_TYPE_BY_STAGE = {
    "rnaseq_qc": "rnaseq",
    "rnaseq_align": "rnaseq",
    "rnaseq_quant": "rnaseq",
    "rnaseq_sample_qc": "rnaseq",
    "rnaseq_de": "rnaseq",
    "rnaseq_enrich": "rnaseq",
    "rnaseq_wgcna": "rnaseq",
    "rnaseq_fgsea": "rnaseq",
    "rnaseq_candidate_family": "rnaseq",
}


def split_md_row(line: str) -> list[str]:
    return [cell.strip() for cell in line.strip().strip("|").split("|")]


def list_expr(text: str, key: str) -> list[str]:
    match = re.search(rf"{re.escape(key)}=\[([^\]]*)\]", text)
    if not match:
        return []
    return [item.strip().strip("'\"` ") for item in re.split(r",|，", match.group(1)) if item.strip().strip("'\"` ")]


def normalize_stage(stage: str) -> str:
    stage = stage.strip().strip("` ")
    aliases = {
        "QC": "rnaseq_qc",
        "align": "rnaseq_align",
        "quant": "rnaseq_quant",
        "DE": "rnaseq_de",
        "evidence_search": "evidence_search",
        "biological_interpretation": "biological_interpretation",
    }
    return aliases.get(stage, stage)


def target_stages(text: str) -> list[str]:
    stages = list_expr(text, "target_stages") or list_expr(text, "target")
    if stages:
        return [normalize_stage(s) for s in stages]
    mapped: list[str] = []
    for stage, terms in STAGE_TERMS.items():
        if any(re.search(rf"\b{re.escape(term)}\b", text) for term in terms):
            mapped.append(stage)
    return list(dict.fromkeys(mapped))


def execution_plan(text: str, expected_stages: list[str]) -> list[str]:
    stages = list_expr(text, "execution_plan")
    if stages:
        return [normalize_stage(s) for s in stages]
    arrow = re.search(r"([A-Za-z0-9_]+(?:\s*->\s*[A-Za-z0-9_]+)+)", text)
    if arrow:
        return [normalize_stage(part.strip()) for part in arrow.group(1).split("->")]
    return expected_stages


def params(text: str) -> dict[str, Any]:
    out: dict[str, Any] = {}
    patterns = {
        "promoter_length": r"promoter_length=([0-9]+)",
        "max_motifs": r"max_motifs=([0-9]+)",
        "bootstrap": r"bootstrap=([0-9]+)",
        "geneFamily": r"family=([A-Za-z0-9_.-]+)",
        "species": r"species=([A-Za-z_]+)",
        "analysis_input": r"analysis_input=([A-Za-z0-9_.-]+)",
    }
    for key, pattern in patterns.items():
        match = re.search(pattern, text)
        if not match:
            continue
        value: Any = match.group(1)
        if key in {"promoter_length", "max_motifs", "bootstrap"}:
            value = int(value)
        out[key] = value
    if "tree_method=ML" in text:
        out["tree_method"] = "ML"
    return out


def entities(text: str, query: str) -> dict[str, Any]:
    combined = text + " " + query
    out: dict[str, Any] = {}
    if "family=" in combined:
        match = re.search(r"family=([A-Za-z0-9_.-]+)", combined)
        if match:
            out["gene_family"] = match.group(1)
    for fam in ["WRKY", "SPL", "NAC", "AP2-ERF", "MYB", "BBX", "P450"]:
        if fam in combined and "gene_family" not in out:
            out["gene_family"] = fam
    for key, value in SPECIES_ALIASES.items():
        if key in combined and "species" not in out:
            out["species"] = value
    if "stress=salt" in combined or "盐" in query:
        out["stress"] = "salt"
    if "id_type=gene_id" in combined:
        out["id_type"] = "gene_id"
    return out


def intent(text: str, layer: str, stages: list[str]) -> str:
    match = re.search(r"intent=`([^`]+)`", text)
    if match:
        return match.group(1)
    if "candidate" in text or "候选" in text:
        return "candidate_evidence_assessment"
    if layer == "evidence_constraint":
        return "evidence_constrained_assessment"
    if layer == "biological_interpretation":
        return "biological_interpretation"
    if layer == "boundary_control":
        return "boundary_control"
    if stages:
        return f"{stages[-1]}_analysis"
    return layer


def required_inputs(text: str) -> list[str]:
    values = list_expr(text, "required_inputs") or list_expr(text, "selected_inputs")
    return list(dict.fromkeys(values))


def expected_outputs(text: str) -> list[str]:
    return list(dict.fromkeys(list_expr(text, "expected_outputs")))


def boundary_action(text: str) -> str:
    match = re.search(r"boundary_action=([A-Za-z0-9_.-]+)", text)
    if match:
        return match.group(1)
    if "can_execute=false" in text:
        return "block_or_request_missing_inputs"
    if "external_or_manual" in text:
        return "mark_external_or_manual"
    if "不能证明" in text or "机制未证明" in text:
        return "claim_limited"
    if "evidence" in text.lower() or "证据" in text:
        return "evidence_supported_or_limited"
    return "none"


def can_execute(text: str) -> bool | str:
    if "can_execute=false" in text:
        return False
    if "can_execute=true" in text:
        return True
    return "depends_on_inputs"


def support_status(text: str) -> str | None:
    for status in ["supported_default", "supported_alternative", "partially_supported", "external_or_manual", "unsupported"]:
        if status in text:
            return status
    return None


def forbidden_stages(must_not: str) -> list[str]:
    found: list[str] = []
    for stage, terms in STAGE_TERMS.items():
        for term in terms:
            if any(f"{action}{term}" in must_not or f"{action} {term}" in must_not for action in ACTION_PATTERNS):
                found.append(stage)
    if "不要安排 expression、motif 或 RNA-seq" in must_not or "不要运行 raw RNA-seq" in must_not:
        found.extend(["rnaseq_qc", "rnaseq_align", "rnaseq_quant", "rnaseq_de", "rnaseq_enrich"])
    return sorted(set(found))


def must_include(text: str) -> list[str]:
    match = re.search(r"must_include.*?\[([^\]]+)\]", text)
    if not match:
        return []
    return [item.strip().strip("'\"` ") for item in re.split(r",|，", match.group(1)) if item.strip().strip("'\"` ")]


def analysis_type(layer: str, stages: list[str]) -> str:
    if layer in {"evidence_constraint", "biological_interpretation", "boundary_control"}:
        return layer
    if any(ANALYSIS_TYPE_BY_STAGE.get(stage) == "rnaseq" for stage in stages):
        return "rnaseq"
    return "gene_family"


def convert(md_path: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    tasks: list[dict[str, Any]] = []
    golds: list[dict[str, Any]] = []
    for line in md_path.read_text(encoding="utf-8").splitlines():
        if not line.startswith("| PGB-"):
            continue
        cells = split_md_row(line)
        if len(cells) < 5:
            continue
        task_id, difficulty, query, expected_text, must_not = cells[:5]
        layer = LAYER_BY_CODE[task_id.split("-")[1]]
        stages = target_stages(expected_text)
        plan = execution_plan(expected_text, stages)
        p = params(expected_text)
        ents = entities(expected_text, query)
        expected = {
            "intent": intent(expected_text, layer, stages),
            "entities": ents,
            "target_stages": stages,
            "execution_plan": plan,
            "parameters": p,
            "artifacts": {
                "required_inputs": required_inputs(expected_text),
                "selected_inputs": list_expr(expected_text, "selected_inputs"),
                "expected_outputs": expected_outputs(expected_text),
            },
            "evidence": [],
            "boundary_action": boundary_action(expected_text),
            "can_execute": can_execute(expected_text),
            "support_status": support_status(expected_text),
            "must_include": must_include(expected_text),
            "must_not_do": [part.strip() for part in re.split(r"；|;", must_not) if part.strip()],
            "rationale": expected_text,
        }
        atype = analysis_type(layer, stages)
        tasks.append({
            "task_id": task_id,
            "layer": layer,
            "category": layer,
            "difficulty": difficulty,
            "source_type": "synthetic_expert_curated",
            "user_query": query,
            "context": {
                "files": [],
                "metadata": {},
                "literature_excerpt": query if layer == "evidence_constraint" else None,
                "result_excerpt": query if layer == "biological_interpretation" else None,
                "manifest_excerpt": None,
            },
            "input_context": {
                "job_id": "",
                "prefix": "",
                "intent": {
                    "analysis_type": atype,
                    "stages": stages,
                    "params": p,
                    "gene_family": ents.get("gene_family"),
                    "species": ents.get("species"),
                },
            },
            "scoring_profile": PROFILE_BY_LAYER[layer],
            "tags": [layer, difficulty],
        })
        golds.append({
            "task_id": task_id,
            "layer": layer,
            "gold": {
                "analysis_type": atype,
                "expected_stages": stages,
                "expected_execution_plan": plan,
                "forbidden_stages": forbidden_stages(must_not),
                "required_params": p,
                "required_dependencies": {},
                **expected,
            },
            "scoring_profile": PROFILE_BY_LAYER[layer],
        })
    return tasks, golds


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")


def run() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--md", type=Path, required=True)
    parser.add_argument("--tasks-out", type=Path, required=True)
    parser.add_argument("--gold-out", type=Path, required=True)
    args = parser.parse_args()
    tasks, golds = convert(args.md)
    write_jsonl(args.tasks_out, tasks)
    write_jsonl(args.gold_out, golds)
    print(json.dumps({"tasks": len(tasks), "gold": len(golds)}, ensure_ascii=False))
    return 0


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