#!/usr/bin/env python3
"""Full PlantGeneBench scorer for five-layer evaluation."""

from __future__ import annotations

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

STAGE_DEPENDENCIES = {
    "phylogeny": ["identify"],
    "motif": ["identify"],
    "domain": ["identify"],
    "structure": ["identify"],
    "promoter": ["identify"],
    "expression": [],
    "synteny": ["identify"],
    "karyotype": ["synteny"],
    "interspecies_synteny": ["identify"],
    "ppi": [],
    "kaks": ["ortholog"],
    "combo": ["phylogeny", "motif", "domain", "structure"],
    "reference_classification": ["phylogeny"],
    "rnaseq_align": ["rnaseq_qc"],
    "rnaseq_quant": ["rnaseq_align"],
    "rnaseq_sample_qc": ["rnaseq_quant"],
    "rnaseq_de": ["rnaseq_quant"],
    "rnaseq_enrich": ["rnaseq_de"],
    "rnaseq_candidate_family": ["rnaseq_de"],
}

LAYER_WEIGHTS = {
    "task_understanding": {"intent": 0.25, "entities": 0.25, "target_stage": 0.25, "boundary": 0.15, "hallucination": 0.10},
    "execution_planning": {"target_stage": 0.20, "execution_plan": 0.25, "parameters": 0.15, "artifacts": 0.15, "boundary": 0.15, "hallucination": 0.10},
    "evidence_constraint": {"target_stage": 0.15, "evidence": 0.40, "boundary": 0.20, "hallucination": 0.25},
    "biological_interpretation": {"evidence": 0.25, "biology": 0.35, "boundary": 0.20, "hallucination": 0.20},
    "boundary_control": {"boundary": 0.45, "hallucination": 0.35, "target_stage": 0.20},
}

HALLUCINATION_PATTERNS = [
    "已完成", "已验证", "证明", "validated", "confirmed", "机制已", "抗盐基因", "完整复现", "已报道",
]

LIMITING_PATTERNS = [
    "候选", "可能", "提示", "需要", "不能", "不足", "验证", "limited", "candidate", "not proven", "manual", "缺少", "不可用",
]


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for lineno, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise SystemExit(f"{path}:{lineno}: invalid JSON: {exc}") from exc
    return rows


def dump_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 as_list(value: Any) -> list[str]:
    if value is None:
        return []
    if isinstance(value, list):
        return [str(v) for v in value]
    return [str(value)]


def norm(value: str) -> str:
    return str(value).strip().lower().replace(" ", "_")


def f1(predicted: list[str], expected: list[str]) -> float:
    p = {norm(x) for x in predicted if x}
    e = {norm(x) for x in expected if x}
    if not p and not e:
        return 1.0
    if not p or not e:
        return 0.0
    hit = len(p & e)
    precision = hit / len(p)
    recall = hit / len(e)
    return 0.0 if precision + recall == 0 else 2 * precision * recall / (precision + recall)


def scalar_equal(predicted: Any, expected: Any) -> bool:
    if isinstance(expected, list):
        return sorted(as_list(predicted)) == sorted(as_list(expected))
    return str(predicted).strip() == str(expected).strip()


def score_params(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    required = gold.get("parameters") or gold.get("required_params") or {}
    if not required:
        return 1.0
    hits = 0
    for key, value in required.items():
        if key in pred and scalar_equal(pred.get(key), value):
            hits += 1
    return hits / len(required)


def score_entities(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    expected = gold.get("entities") or {}
    if not expected:
        return 1.0
    predicted = pred.get("entities") or {}
    params = pred.get("params") or {}
    aliases = {
        "gene_family": ["gene_family", "geneFamily", "family", "target_family"],
        "species": ["species", "organism", "target_species"],
        "tissue": ["tissue", "sample", "condition"],
    }
    hits = 0
    total = 0
    for key, value in expected.items():
        if value in (None, "", []):
            continue
        total += 1
        got = None
        for candidate in aliases.get(key, [key]):
            got = predicted.get(candidate) or params.get(candidate)
            if got not in (None, "", []):
                break
        if got and norm(got) == norm(value):
            hits += 1
    return 1.0 if total == 0 else hits / total


def score_intent(pred: dict[str, Any], gold: dict[str, Any], target_stage_score: float) -> float:
    expected = gold.get("intent")
    if not expected:
        return 1.0
    predicted = pred.get("intent")
    if predicted and norm(predicted) == norm(expected):
        return 1.0
    analysis_ok = not gold.get("analysis_type") or norm(pred.get("analysis_type") or "") == norm(gold.get("analysis_type") or "")
    if analysis_ok and target_stage_score >= 0.99:
        return 1.0
    if analysis_ok and target_stage_score >= 0.5:
        return 0.5
    return 0.0

def score_execution_plan(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    expected = as_list(gold.get("execution_plan") or gold.get("expected_execution_plan") or gold.get("target_stages"))
    predicted = as_list(pred.get("execution_plan") or pred.get("stages"))
    return f1(predicted, expected)


def score_artifacts(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    expected = ((gold.get("artifacts") or {}).get("required_inputs") or []) + ((gold.get("artifacts") or {}).get("expected_outputs") or [])
    if not expected:
        return 1.0
    predicted_artifacts = pred.get("artifacts") or {}
    predicted = []
    if isinstance(predicted_artifacts, dict):
        for value in predicted_artifacts.values():
            if isinstance(value, dict):
                for nested in value.values():
                    predicted.extend(as_list(nested))
            else:
                predicted.extend(as_list(value))
    answer = str(pred.get("answer_text") or "")
    predicted.extend([item for item in expected if item in answer])
    return f1(predicted, expected)


def has_hallucination(pred: dict[str, Any], gold: dict[str, Any]) -> bool:
    text = str(pred.get("answer_text") or "")
    must_not = " ".join(as_list(gold.get("must_not_do")))
    if not text:
        return False
    if any(pattern in text for pattern in HALLUCINATION_PATTERNS) and any(trigger in must_not for trigger in ["不要声称", "不要写成", "不要编造", "不要标记", "不要输出不存在", "不要生成"]):
        return True
    return False


def score_hallucination(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    return 0.0 if has_hallucination(pred, gold) else 1.0


def expected_boundary_positive(gold: dict[str, Any]) -> bool:
    action = str(gold.get("boundary_action") or "")
    can_execute = gold.get("can_execute")
    must_not = " ".join(as_list(gold.get("must_not_do")))
    return bool(action and action != "none") or can_execute is False or "不要" in must_not


def score_boundary(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    if not expected_boundary_positive(gold):
        return 1.0
    text = " ".join([str(pred.get("answer_text") or ""), str(pred.get("backend_status") or ""), str(pred.get("boundary_action") or "")])
    if any(pattern in text for pattern in LIMITING_PATTERNS):
        return 1.0
    if pred.get("backend_status") in {"missing_input_blocked", "backend_success_false"} and gold.get("can_execute") is False:
        return 1.0
    return 0.0


def evidence_has_content(evidence: Any) -> bool:
    if not evidence:
        return False
    if isinstance(evidence, list):
        for item in evidence:
            if isinstance(item, dict):
                if any(str(v).strip() for v in item.values()):
                    return True
            elif str(item).strip():
                return True
        return False
    return bool(str(evidence).strip())


def score_evidence(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    expected_status = gold.get("support_status")
    text = str(pred.get("answer_text") or "")
    evidence = pred.get("evidence") or []
    if expected_status == "external_or_manual":
        return 1.0 if ("manual" in text.lower() or "人工" in text or "外部" in text or "qRT-PCR" in text) else 0.5 if pred.get("backend_status") in {"missing_input_blocked", "backend_success_false"} else 0.0
    if expected_status and expected_status in text:
        return 1.0
    if evidence_has_content(evidence):
        return 1.0
    if pred.get("backend_status") in {"missing_input_blocked", "backend_success_false"}:
        return 0.5
    # For plan-only outputs, absence of evidence is not automatically a hallucination but needs manual review.
    return 0.5


def score_biology(pred: dict[str, Any], gold: dict[str, Any]) -> float:
    text = str(pred.get("answer_text") or "")
    if not text:
        return 0.5 if pred.get("backend_status") in {"missing_input_blocked", "backend_success_false"} else 0.0
    if has_hallucination(pred, gold):
        return 0.0
    if any(pattern in text for pattern in LIMITING_PATTERNS):
        return 1.0
    return 0.5


def classify_errors(components: dict[str, float], pred: dict[str, Any], gold: dict[str, Any]) -> list[str]:
    errors = []
    for key, value in components.items():
        if key in {"stage_precision", "stage_recall"}:
            continue
        if value < 1.0:
            errors.append(f"{key}_partial" if value > 0 else f"{key}_fail")
    status = pred.get("backend_status")
    if status and status != "success":
        errors.append(status)
    if pred.get("mode") == "deepseek_json" and pred.get("llm_json_valid") is False:
        errors.append("llm_json_invalid")
    if has_hallucination(pred, gold):
        errors.append("hallucination_detected")
    return errors


def score_one(pred_row: dict[str, Any], gold_row: dict[str, Any]) -> dict[str, Any]:
    pred = pred_row.get("prediction") or {}
    gold = gold_row.get("gold") or {}
    layer = gold_row.get("layer") or pred_row.get("request", {}).get("layer") or "unknown"
    expected_stages = as_list(gold.get("target_stages") or gold.get("expected_stages"))
    predicted_stages = as_list(pred.get("stages"))
    target_stage_score = f1(predicted_stages, expected_stages)
    components = {
        "intent": score_intent(pred, gold, target_stage_score),
        "entities": score_entities(pred, gold),
        "target_stage": target_stage_score,
        "execution_plan": score_execution_plan(pred, gold),
        "parameters": score_params(pred.get("params") or {}, gold),
        "artifacts": score_artifacts(pred, gold),
        "evidence": score_evidence(pred, gold),
        "biology": score_biology(pred, gold),
        "boundary": score_boundary(pred, gold),
        "hallucination": score_hallucination(pred, gold),
    }
    weights = LAYER_WEIGHTS.get(layer, {"target_stage": 1.0})
    total_weight = sum(weights.values()) or 1.0
    score = sum(components.get(key, 0.0) * weight for key, weight in weights.items()) / total_weight
    errors = classify_errors(components, pred, gold)
    provider = pred_row.get("provider", "")
    mode = pred.get("mode", "")
    review_required = (
        layer in {"evidence_constraint", "biological_interpretation", "boundary_control"}
        or score < 0.8
        or provider == "deepseek_text"
        or (mode == "deepseek_json" and pred.get("llm_json_valid") is not True)
        or any(e in errors for e in ["hallucination_detected", "backend_exception", "http_or_json_error", "llm_json_invalid"])
    )
    return {
        "task_id": pred_row["task_id"],
        "provider": provider,
        "run_id": pred_row.get("run_id", ""),
        "layer": layer,
        "score": round(score, 6),
        "components": {k: round(v, 6) for k, v in components.items()},
        "weights": weights,
        "expected_stages": expected_stages,
        "predicted_stages": predicted_stages,
        "backend_status": pred.get("backend_status", ""),
        "errors": errors,
        "review_required": review_required,
        "prediction": pred,
        "gold": gold,
        "llm_json_valid": pred.get("llm_json_valid"),
        "llm_json_status": pred.get("llm_json_status", ""),
        "llm_model": pred.get("llm_model", ""),
        "llm_usage": pred.get("llm_usage") or {},
    }


def write_manual_review(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = ["task_id", "provider", "layer", "machine_score", "backend_status", "errors", "llm_json_valid", "llm_json_status", "llm_model", "manual_score", "raw_answer_supported", "json_schema_issue", "evidence_faithfulness", "biology_reasoning", "boundary_correct", "llm_overclaim", "evidence_citation_status", "hallucination_type", "reviewer_note", "final_decision"]
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t")
        writer.writeheader()
        for row in rows:
            if not row.get("review_required"):
                continue
            writer.writerow({
                "task_id": row["task_id"],
                "provider": row.get("provider", ""),
                "layer": row.get("layer", ""),
                "machine_score": row.get("score", 0),
                "backend_status": row.get("backend_status", ""),
                "errors": ";".join(row.get("errors") or []),
                "llm_json_valid": row.get("llm_json_valid", ""),
                "llm_json_status": row.get("llm_json_status", ""),
                "llm_model": row.get("llm_model", ""),
                "manual_score": "",
                "raw_answer_supported": "",
                "json_schema_issue": "",
                "evidence_faithfulness": "",
                "biology_reasoning": "",
                "boundary_correct": "",
                "llm_overclaim": "",
                "evidence_citation_status": "",
                "hallucination_type": "",
                "reviewer_note": "",
                "final_decision": "",
            })


def run() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--predictions", type=Path, required=True)
    parser.add_argument("--gold", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--manual-review-out", type=Path, default=None)
    args = parser.parse_args()
    predictions = load_jsonl(args.predictions)
    gold_by_id = {row["task_id"]: row for row in load_jsonl(args.gold)}
    scored = []
    for pred in predictions:
        gold = gold_by_id.get(pred.get("task_id"))
        if not gold:
            continue
        scored.append(score_one(pred, gold))
    dump_jsonl(args.out, scored)
    if args.manual_review_out:
        write_manual_review(args.manual_review_out, scored)
    print(f"Wrote {len(scored)} scored rows to {args.out}")
    if args.manual_review_out:
        print(f"Wrote manual review queue to {args.manual_review_out}")
    return 0


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