#!/usr/bin/env python3
"""Aggregate MCP scaling experiment metrics from Biomni/bioagent-bench run dirs."""

from __future__ import annotations

import argparse
import json
import re
import time
from functools import lru_cache
from pathlib import Path
from typing import Any

DEFAULT_EXP_DIR = Path(__file__).resolve().parents[1]
DEFAULT_GOLD = DEFAULT_EXP_DIR / "gold_tools.json"
DEFAULT_MANIFEST = DEFAULT_EXP_DIR / "configs" / "mcp_scale_manifest.json"
DEFAULT_DATASET_ROOT = Path("/225040511/project/bioagent-bench/dataset")
DEFAULT_PRIMARY_K = 10

TYPE_MISMATCH_PATTERNS = [
    r"typeerror",
    r"valueerror",
    r"validationerror",
    r"pydantic",
    r"missing required positional argument",
    r"unexpected keyword argument",
    r"unsupported operand type",
    r"input should be",
    r"should be of type",
    r"must be (?:a|an)?\s*(?:int|integer|float|bool|boolean|string|str|list|dict|path|file)",
    r"expected .*?(?:int|integer|float|bool|boolean|string|str|list|dict|path|file)",
    r"invalid literal for int",
    r"cannot be interpreted as an integer",
    r"is not iterable",
]

CONSTRAINT_VIOLATION_PATTERNS = [
    r"\bpip install\b",
    r"\bconda install\b",
    r"\bapt-get\b",
    r"\bwget\b",
    r"\bcurl\b",
    r"\bgit clone\b",
    r"https?://",
]


def load_json(path: Path, default: Any = None) -> Any:
    if not path.exists():
        return default
    return json.loads(path.read_text(encoding="utf-8"))


def load_text(path: Path, default: str = "") -> str:
    if not path.exists():
        return default
    return path.read_text(encoding="utf-8", errors="ignore")


def load_yaml_if_available(path: Path, default: Any = None) -> Any:
    if not path.exists():
        return default
    try:
        import yaml
    except ModuleNotFoundError:
        return default
    return yaml.safe_load(path.read_text(encoding="utf-8")) or default


def normalize(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", text.lower())


def extract_names(obj: Any) -> list[str]:
    names: list[str] = []
    if isinstance(obj, dict):
        for key, value in obj.items():
            if key in {"name", "tool_name", "server", "module"} and isinstance(value, str):
                names.append(value)
            names.extend(extract_names(value))
    elif isinstance(obj, list):
        for item in obj:
            names.extend(extract_names(item))
    elif isinstance(obj, str):
        for token in re.findall(r"\b[a-zA-Z][a-zA-Z0-9_.-]{2,}\b", obj):
            if any(
                hint in token.lower()
                for hint in (
                    "mcp",
                    "tool",
                    "kraken",
                    "kallisto",
                    "salmon",
                    "scanpy",
                    "star",
                    "bwa",
                    "bcftools",
                    "samtools",
                    "deseq",
                    "kaiju",
                    "megahit",
                    "gatk",
                    "freebayes",
                )
            ):
                names.append(token)
    return names


def recall_at_k(candidates: list[str], gold: list[str], k: int) -> float:
    if not gold:
        return 0.0
    top = {normalize(x) for x in candidates[:k]}
    hits = 0
    for item in gold:
        n = normalize(item)
        if any(n in candidate or candidate in n for candidate in top):
            hits += 1
    return hits / len(gold)


def gold_hit(candidates: list[str], gold: list[str]) -> bool:
    return recall_at_k(candidates, gold, max(len(candidates), 1)) > 0


def mean_or_none(values: list[float | int | None]) -> float | None:
    numeric = [float(value) for value in values if isinstance(value, (int, float))]
    if not numeric:
        return None
    return sum(numeric) / len(numeric)


def workflow_validity(evidence: list[str], workflow: list[dict[str, Any]]) -> dict[str, Any]:
    found = []
    cursor = 0
    normalized_evidence = [normalize(x) for x in evidence]
    for step in workflow:
        patterns = step.get("tool_patterns", [])
        matched_at = None
        for i in range(cursor, len(normalized_evidence)):
            row = normalized_evidence[i]
            if any(normalize(pattern) in row for pattern in patterns):
                matched_at = i
                break
        if matched_at is not None:
            found.append(step.get("name", "step"))
            cursor = matched_at + 1
    return {
        "valid": len(found) == len(workflow) and bool(workflow),
        "steps_matched": len(found),
        "steps_total": len(workflow),
        "matched_step_names": found,
    }


def execution_success(run_dir: Path) -> bool:
    validation = load_json(run_dir / "output_validation.json", {})
    if isinstance(validation, dict):
        if validation.get("success") is True or validation.get("all_outputs_exist") is True:
            return True
        outputs = validation.get("outputs") or validation.get("output_paths")
        if isinstance(outputs, list) and outputs:
            return all(bool(x.get("exists")) for x in outputs if isinstance(x, dict))
    summary = load_json(run_dir / "run_summary.json", {})
    if isinstance(summary, dict):
        if summary.get("success") is True:
            return True
        if summary.get("status") in {"success", "completed"}:
            return True
    metadata = load_json(run_dir / "run_metadata.json", {})
    for output in metadata.get("output_paths", []) if isinstance(metadata, dict) else []:
        if not Path(output).exists():
            return False
    return bool(metadata.get("output_paths")) if isinstance(metadata, dict) else False


def rough_token_count(text: str) -> int:
    return max(1, len(re.findall(r"\S+", text)))


def discover_task_ids(dataset_root: Path) -> list[str]:
    if not dataset_root.exists():
        return []
    return sorted(p.name for p in dataset_root.iterdir() if p.is_dir())


def collect_task_runs(runs_root: Path, task_id: str) -> list[Path]:
    if not runs_root.exists():
        return []
    pattern = re.compile(rf"^{re.escape(task_id)}_(\d{{8}}_\d{{6}})$")
    matched: list[tuple[str, Path]] = []
    for run_dir in runs_root.iterdir():
        if not run_dir.is_dir():
            continue
        match = pattern.match(run_dir.name)
        if match:
            matched.append((match.group(1), run_dir))
    matched.sort(key=lambda x: x[0])
    return [path for _, path in matched]


@lru_cache(maxsize=16)
def load_mcp_tool_universe(config_path: str) -> list[str]:
    cfg = load_yaml_if_available(Path(config_path), default={})
    if not isinstance(cfg, dict):
        return []
    names: list[str] = []
    for server_meta in (cfg.get("mcp_servers") or {}).values():
        if not isinstance(server_meta, dict):
            continue
        for tool in server_meta.get("tools", []):
            if isinstance(tool, dict) and tool.get("name"):
                names.append(tool["name"])
    return names


def selected_tool_names(retrieval_plan: dict[str, Any]) -> list[str]:
    names = retrieval_plan.get("selected_resource_names", {}).get("tools", [])
    if names:
        return [str(name) for name in names if name]
    tool_items = retrieval_plan.get("selected_resources", {}).get("tools", [])
    return [item.get("name") for item in tool_items if isinstance(item, dict) and item.get("name")]


def registered_tool_names(retrieval_plan: dict[str, Any]) -> list[str]:
    return [str(name) for name in retrieval_plan.get("registered_tool_names", []) if name]


def stored_selected_tool_count(retrieval_plan: dict[str, Any]) -> int | None:
    count = retrieval_plan.get("selected_tool_count")
    return count if isinstance(count, int) else None


def retrieval_metadata_available(retrieval_plan: dict[str, Any]) -> bool:
    if selected_tool_names(retrieval_plan):
        return True
    if registered_tool_names(retrieval_plan):
        return True
    planning_context = retrieval_plan.get("planning_context_text")
    if isinstance(planning_context, str) and planning_context.strip():
        return True
    query_context = retrieval_plan.get("query_context")
    if isinstance(query_context, dict) and query_context:
        return True
    return False


def retrieval_candidates_from_plan(retrieval_plan: dict[str, Any]) -> list[str]:
    candidates: list[str] = []
    candidates.extend(selected_tool_names(retrieval_plan))
    candidates.extend(registered_tool_names(retrieval_plan))
    planning_context = retrieval_plan.get("planning_context_text")
    if isinstance(planning_context, str) and planning_context.strip():
        candidates.extend(extract_names(planning_context))
    query_context = retrieval_plan.get("query_context")
    if isinstance(query_context, dict) and query_context:
        candidates.extend(extract_names(query_context))
    return sorted(dict.fromkeys(name for name in candidates if name))


def planning_context_fallback(retrieval_plan: dict[str, Any], metadata: dict[str, Any], run_dir: Path) -> str:
    planning_context_text = retrieval_plan.get("planning_context_text")
    if isinstance(planning_context_text, str) and planning_context_text.strip():
        return planning_context_text

    query_context = retrieval_plan.get("query_context")
    if isinstance(query_context, dict) and query_context:
        return json.dumps(query_context, ensure_ascii=False, default=str)

    query = retrieval_plan.get("query")
    if isinstance(query, str) and query.strip():
        return query

    metadata_query = metadata.get("query") if isinstance(metadata, dict) else None
    if isinstance(metadata_query, str) and metadata_query.strip():
        return metadata_query

    task_query = load_text(run_dir / "task_query.txt").strip()
    if task_query:
        return task_query

    return ""


def non_human_log_text(execution_log: dict[str, Any], execution_text_path: Path, final_text: str) -> str:
    entries = execution_log.get("log_entries", []) if isinstance(execution_log, dict) else []
    filtered = []
    for entry in entries:
        text = str(entry)
        if "================================ Human Message" in text:
            continue
        filtered.append(text)
    if filtered:
        return "\n\n".join(filtered) + "\n" + final_text
    if execution_text_path.exists():
        return execution_text_path.read_text(encoding="utf-8", errors="ignore") + "\n" + final_text
    return final_text


def matched_tool_names(text: str, candidate_names: list[str]) -> list[str]:
    haystack = normalize(text)
    matched = []
    for name in candidate_names:
        token = normalize(name)
        if token and token in haystack:
            matched.append(name)
    return sorted(set(matched))


def execution_attempt_count(text: str, executed_tools: list[str]) -> int:
    counts = [
        len(re.findall(r"<execute>", text, flags=re.IGNORECASE)),
        len(re.findall(r"command_executed", text, flags=re.IGNORECASE)),
        len(executed_tools),
    ]
    return max(1, *counts)


def pattern_rate(text: str, patterns: list[str], denominator: int) -> tuple[float, int]:
    count = 0
    for pattern in patterns:
        count += len(re.findall(pattern, text, flags=re.IGNORECASE))
    return min(1.0, count / max(1, denominator)), count


def biological_constraint_rate(text: str, metadata: dict[str, Any], denominator: int) -> tuple[float, int]:
    count = 0
    for pattern in CONSTRAINT_VIOLATION_PATTERNS:
        count += len(re.findall(pattern, text, flags=re.IGNORECASE))

    task_id = metadata.get("task_id", "") if isinstance(metadata, dict) else ""
    if task_id:
        sibling_pattern = rf"/225040511/project/bioagent-bench/dataset/(?!{re.escape(task_id)})([a-zA-Z0-9_.-]+)"
        count += len(re.findall(sibling_pattern, text))
    count += len(re.findall(r"/225040511/project/bioagent-bench/dataset/.+?/results", text))
    return min(1.0, count / max(1, denominator)), count


def summarize_run(run_dir: Path, gold: dict[str, Any], ks: list[int], primary_k: int) -> dict[str, Any]:
    metadata = load_json(run_dir / "run_metadata.json", {})
    task_id = metadata.get("task_id") if isinstance(metadata, dict) else None
    if not task_id:
        task_id = re.sub(r"_20\d{6}_\d{6}$", "", run_dir.name)
    gold_entry = gold.get(task_id, {})

    retrieval_plan = load_json(run_dir / "retrieval_plan.json", {})
    execution_log = load_json(run_dir / "execution_log.json", {})
    execution_text_path = run_dir / "execution_log.txt"
    final_text = (run_dir / "final_answer.txt").read_text(encoding="utf-8", errors="ignore") if (run_dir / "final_answer.txt").exists() else ""

    retrieval_candidates = retrieval_candidates_from_plan(retrieval_plan)
    retrieval_available = retrieval_metadata_available(retrieval_plan)
    filtered_execution_text = non_human_log_text(execution_log, execution_text_path, final_text)
    execution_evidence = extract_names(execution_log) + re.findall(
        r"\b[a-zA-Z][a-zA-Z0-9_.-]{2,}\b",
        filtered_execution_text,
    )

    gold_tools = gold_entry.get("gold_tools", [])
    gold_servers = gold_entry.get("gold_servers", [])
    all_gold = [*gold_tools, *gold_servers]
    recall = {
        str(k): (recall_at_k(retrieval_candidates, all_gold, k) if retrieval_available else None)
        for k in ks
    }

    planning_context_text = planning_context_fallback(retrieval_plan, metadata, run_dir)

    timing = metadata.get("timing", {}) if isinstance(metadata, dict) else {}
    metrics = metadata.get("metrics", {}) if isinstance(metadata, dict) else {}
    planning_latency = (
        retrieval_plan.get("planning_latency_seconds")
        or timing.get("planning_latency_seconds")
        or metrics.get("planning_latency")
        or metrics.get("plan_seconds")
    )

    mcp_config_path = retrieval_plan.get("mcp_config") or metadata.get("mcp_config")
    mcp_tool_universe = load_mcp_tool_universe(str(mcp_config_path)) if mcp_config_path else []
    selected_tools = selected_tool_names(retrieval_plan)
    registered_tools = registered_tool_names(retrieval_plan)
    retrieved_tool_candidates = set(selected_tools or registered_tools)
    executed_mcp_tools = matched_tool_names(filtered_execution_text, mcp_tool_universe) if mcp_tool_universe else []
    hallucinated_tools = sorted(set(executed_mcp_tools) - retrieved_tool_candidates)
    hallucinated_tool_rate = (
        len(hallucinated_tools) / len(executed_mcp_tools)
        if executed_mcp_tools and retrieved_tool_candidates
        else None
    )

    attempts = execution_attempt_count(filtered_execution_text, executed_mcp_tools)
    data_type_mismatch_rate, data_type_mismatch_events = pattern_rate(filtered_execution_text, TYPE_MISMATCH_PATTERNS, attempts)
    biological_constraint_error_rate, biological_constraint_events = biological_constraint_rate(
        filtered_execution_text,
        metadata if isinstance(metadata, dict) else {},
        attempts,
    )

    workflow = workflow_validity(execution_evidence, gold_entry.get("workflow", []))
    execution_ok = execution_success(run_dir)

    return {
        "task_id": task_id,
        "run_dir": str(run_dir),
        "scale_label": run_dir.parent.name,
        "retrieval_metadata_available": retrieval_available,
        "retrieval_recall_at_k": recall,
        "retrieval_recall_at_primary_k": recall.get(str(primary_k)),
        "retrieval_recall_k": primary_k,
        "tool_selection_accuracy": gold_hit(execution_evidence, all_gold),
        "workflow_validity": workflow,
        "execution_success": execution_ok,
        "context_tokens_rough": rough_token_count(planning_context_text),
        "planning_latency_seconds": planning_latency,
        "retrieved_candidate_count": len(retrieval_candidates),
        "execution_evidence_count": len(execution_evidence),
        "selected_tool_count": len(selected_tools) if selected_tools else stored_selected_tool_count(retrieval_plan) or 0,
        "registered_tool_count": len(registered_tools),
        "executed_mcp_tools": executed_mcp_tools,
        "hallucinated_tools": hallucinated_tools,
        "hallucinated_tool_rate": hallucinated_tool_rate,
        "data_type_mismatch_rate": data_type_mismatch_rate,
        "data_type_mismatch_events": data_type_mismatch_events,
        "biological_constraint_error_rate": biological_constraint_error_rate,
        "biological_constraint_error_events": biological_constraint_events,
        "execution_attempts_estimate": attempts,
    }


def aggregate_results(results: list[dict[str, Any]], ks: list[int], primary_k: int) -> dict[str, Any]:
    aggregate: dict[str, Any] = {"run_count": len(results), "retrieval_recall_k": primary_k}
    if not results:
        return aggregate

    for k in ks:
        vals = [r["retrieval_recall_at_k"][str(k)] for r in results]
        aggregate[f"mean_retrieval_recall@{k}"] = mean_or_none(vals)
    aggregate["tool_selection_accuracy"] = sum(r["tool_selection_accuracy"] for r in results) / len(results)
    aggregate["workflow_validity"] = sum(r["workflow_validity"]["valid"] for r in results) / len(results)
    aggregate["execution_success_rate"] = sum(r["execution_success"] for r in results) / len(results)
    aggregate["mean_context_tokens_rough"] = sum(r["context_tokens_rough"] for r in results) / len(results)
    latencies = [r["planning_latency_seconds"] for r in results if isinstance(r["planning_latency_seconds"], (int, float))]
    aggregate["mean_planning_latency_seconds"] = sum(latencies) / len(latencies) if latencies else None
    aggregate["retrieval_metadata_coverage"] = (
        sum(r["retrieval_metadata_available"] for r in results) / len(results)
    )
    aggregate["mean_hallucinated_tool_rate"] = mean_or_none([r["hallucinated_tool_rate"] for r in results])
    aggregate["mean_data_type_mismatch_rate"] = mean_or_none([r["data_type_mismatch_rate"] for r in results])
    aggregate["mean_biological_constraint_error_rate"] = mean_or_none(
        [r["biological_constraint_error_rate"] for r in results]
    )

    aggregate["Retrieval Recall@k"] = aggregate.get(f"mean_retrieval_recall@{primary_k}")
    aggregate["Workflow Validity"] = aggregate["workflow_validity"]
    aggregate["Execution Success Rate"] = aggregate["execution_success_rate"]
    aggregate["Context Tokens"] = aggregate["mean_context_tokens_rough"]
    aggregate["Planning Latency"] = aggregate["mean_planning_latency_seconds"]
    aggregate["Hallucinated Tool Rate"] = aggregate["mean_hallucinated_tool_rate"]
    aggregate["Data-Type Mismatch Rate"] = aggregate["mean_data_type_mismatch_rate"]
    aggregate["Biological Constraint Error Rate"] = aggregate["mean_biological_constraint_error_rate"]
    return aggregate


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--runs-root", type=Path, required=True, help="Directory containing tier/task run dirs.")
    parser.add_argument(
        "--dataset-root",
        type=Path,
        default=DEFAULT_DATASET_ROOT,
        help="Benchmark dataset root used to enumerate task IDs for per-task outputs.",
    )
    parser.add_argument("--gold", type=Path, default=DEFAULT_GOLD)
    parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    parser.add_argument("--out", type=Path, default=DEFAULT_EXP_DIR / "results" / "metrics_summary.json")
    parser.add_argument(
        "--per-task-dir",
        type=Path,
        default=DEFAULT_EXP_DIR / "results" / "by_task",
        help="Directory where one metrics JSON per task is written.",
    )
    parser.add_argument("--k", type=int, nargs="+", default=[1, 3, 5, 10, 20, 50])
    parser.add_argument("--primary-k", type=int, default=DEFAULT_PRIMARY_K)
    args = parser.parse_args()

    started = time.time()
    gold = load_json(args.gold, {})
    manifest = load_json(args.manifest, {})
    task_ids = discover_task_ids(args.dataset_root)
    args.per_task_dir.mkdir(parents=True, exist_ok=True)

    results: list[dict[str, Any]] = []
    task_index = []
    for task_id in task_ids:
        task_run_dirs = collect_task_runs(args.runs_root, task_id)
        task_results = [summarize_run(run_dir, gold, args.k, args.primary_k) for run_dir in task_run_dirs]
        results.extend(task_results)
        task_payload = {
            "evaluated_at_unix": started,
            "task_id": task_id,
            "runs_root": str(args.runs_root),
            "dataset_root": str(args.dataset_root),
            "manifest": manifest,
            "aggregate": aggregate_results(task_results, args.k, args.primary_k),
            "results": task_results,
        }
        task_out = args.per_task_dir / f"{task_id}.json"
        task_out.write_text(json.dumps(task_payload, indent=2, ensure_ascii=False), encoding="utf-8")
        task_index.append({"task_id": task_id, "run_count": len(task_results), "output": str(task_out)})

    aggregate = aggregate_results(results, args.k, args.primary_k)
    payload = {
        "evaluated_at_unix": started,
        "runs_root": str(args.runs_root),
        "dataset_root": str(args.dataset_root),
        "manifest": manifest,
        "tasks": task_index,
        "aggregate": aggregate,
        "results": results,
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
    print(json.dumps(payload["aggregate"], indent=2, ensure_ascii=False))
    return 0


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