molgate_ui/creadrug_export.py

"""
Export a MolGate / DrugOps session into a CreaDrug24 case pack.

Writes under CREADRUG_ROOT/input/cases/<caseId>/ and updates input/cases_molgate.json.
"""

from __future__ import annotations

import json
import re
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from molgate_ui.index_status import safe_session_dir, safe_session_file
from molgate_ui.paths import SESSIONS, creadrug24_root

CASES_REGISTRY = "input/cases_molgate.json"


def _read_json(p: Path) -> dict[str, Any] | None:
    if not p.is_file():
        return None
    try:
        return json.loads(p.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _first_existing(sd: Path, rel_paths: list[str]) -> Path | None:
    for rel in rel_paths:
        fp = safe_session_file(sd, rel)
        if fp:
            return fp
    return None


def _resolve_drug_id(sd: Path, meta: dict[str, Any] | None, s3: dict[str, Any] | None) -> int | None:
    for src in (meta, s3):
        if not src:
            continue
        raw = src.get("drug_id") or src.get("_catalog_drug_id")
        if raw is not None:
            try:
                return int(raw)
            except (TypeError, ValueError):
                pass
    man = _read_json(sd / "molgate_run_manifest.json")
    if man:
        raw = man.get("drug_id")
        if raw is not None:
            try:
                return int(raw)
            except (TypeError, ValueError):
                pass
    return None


def _drug_name(meta: dict[str, Any] | None, s3: dict[str, Any] | None, dr: dict[str, Any] | None) -> str:
    for src in (meta, s3, dr):
        if not src:
            continue
        for key in ("drug_label", "drug_name", "target_drug", "common_name"):
            val = src.get(key)
            if isinstance(val, str) and val.strip():
                return val.strip()
    return ""


def _target_label(meta: dict[str, Any] | None, s3: dict[str, Any] | None, dr: dict[str, Any] | None) -> str:
    target = str(
        (meta or {}).get("target")
        or (s3 or {}).get("target")
        or (dr or {}).get("target")
        or ""
    ).strip().upper()
    if target.endswith("_PROXY"):
        return target.split("_")[0]
    parts = target.split("_")
    return parts[0] if parts else target


def _default_case_id(session_id: str, drug_id: int | None) -> str:
    if drug_id is not None and drug_id > 0:
        return f"M{drug_id:02d}" if drug_id < 100 else f"M{drug_id}"
    sid = re.sub(r"[^a-fA-F0-9]", "", session_id)[:8]
    return f"M_{sid or 'session'}"


def _sanitize_case_id(case_id: str) -> str:
    raw = (case_id or "").strip().upper()
    if not raw:
        raise ValueError("case_id 不可為空")
    if not re.match(r"^[A-Z][A-Z0-9_\-]{0,31}$", raw):
        raise ValueError(f"case_id 格式不合法：{case_id}")
    return raw


def _build_docking_log(dr: dict[str, Any] | None, vina_log: Path | None) -> list[str]:
    lines: list[str] = []
    if vina_log and vina_log.is_file():
        try:
            tail = vina_log.read_text(encoding="utf-8", errors="replace").splitlines()[-12:]
            lines.extend(t.strip() for t in tail if t.strip())
        except OSError:
            pass
    if dr:
        aff = dr.get("best_affinity")
        mode = dr.get("best_pose_mode", 1)
        tier = dr.get("confidence_tier", "")
        qc = dr.get("qc_pass")
        if aff is not None:
            lines.insert(0, f"Mode {mode}: {aff} kcal/mol")
        if tier:
            lines.append(f"[STAGE 26] Confidence tier {tier}")
        if qc is not None:
            lines.append(f"[STAGE 25] QC {'PASS' if qc else 'WARN'}")
    return lines or ["[DrugOps export] docking log unavailable"]


def _build_summary(dr: dict[str, Any] | None, qc: dict[str, Any] | None) -> str:
    if dr:
        flags = dr.get("flags") or []
        flag_txt = f" Flags: {'; '.join(flags)}" if flags else ""
        return (
            f"DrugOps export · best pose #{dr.get('best_pose_mode', 1)} · "
            f"affinity {dr.get('best_affinity')} kcal/mol · "
            f"tier {dr.get('confidence_tier', '—')}.{flag_txt}"
        )
    if qc:
        return f"DrugOps export · QC status {qc.get('status', '—')}"
    return "DrugOps session export."


def build_case_entry(
    case_id: str,
    case_dir_rel: str,
    *,
    drug_id: int | None,
    drug_name: str,
    smiles: str,
    target_name: str,
    index_key: str,
    pdb_id: str,
    dr: dict[str, Any] | None,
    qc: dict[str, Any] | None,
    vina_log: Path | None,
    pose_ext: str,
) -> dict[str, Any]:
    aff = dr.get("best_affinity") if dr else None
    tier = str(dr.get("confidence_tier") or "—") if dr else "—"
    qc_pass = dr.get("qc_pass") if dr else None
    status = "COMPLETE" if aff is not None else "PARTIAL"
    if qc and str(qc.get("status", "")).upper() in ("FAIL", "ERROR"):
        status = "QC-FLAG"
    return {
        "caseId": case_id,
        "catalogDrugId": drug_id,
        "drugName": drug_name or index_key or case_id,
        "formula": "",
        "smiles": smiles or "",
        "targetName": target_name or index_key,
        "indexKey": index_key,
        "pdbId": pdb_id,
        "receptorPdbUrl": f"{case_dir_rel}/receptor.pdb",
        "dockedPdbqtUrl": f"{case_dir_rel}/best_pose.{pose_ext}",
        "pocketConfigUrl": f"{case_dir_rel}/stage_23_pocket_center.json",
        "dockingResultUrl": f"{case_dir_rel}/docking_result.json",
        "bestAffinity": aff,
        "confidenceTier": tier,
        "status": status if qc_pass is not False else "QC-FLAG",
        "summary": _build_summary(dr, qc),
        "dockingLog": _build_docking_log(dr, vina_log),
        "source": "drugops",
        "exportedAt": datetime.now(timezone.utc).isoformat(),
    }


def _update_registry(creadrug_root: Path, entry: dict[str, Any]) -> Path:
    reg_path = creadrug_root / CASES_REGISTRY
    reg_path.parent.mkdir(parents=True, exist_ok=True)
    payload: dict[str, Any] = {"schema_version": "1.0", "cases": []}
    if reg_path.is_file():
        try:
            existing = json.loads(reg_path.read_text(encoding="utf-8"))
            if isinstance(existing, dict) and isinstance(existing.get("cases"), list):
                payload = existing
            elif isinstance(existing, list):
                payload = {"schema_version": "1.0", "cases": existing}
        except (json.JSONDecodeError, OSError):
            pass
    cases: list[dict[str, Any]] = payload.setdefault("cases", [])
    cid = entry["caseId"]
    replaced = False
    for i, c in enumerate(cases):
        if c.get("caseId") == cid:
            cases[i] = entry
            replaced = True
            break
    if not replaced:
        cases.append(entry)
    payload["updated_at"] = datetime.now(timezone.utc).isoformat()
    reg_path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    return reg_path


def export_session_to_creadrug(
    session_id: str,
    *,
    sessions_root: Path | None = None,
    creadrug_root: Path | None = None,
    case_id: str | None = None,
    overwrite: bool = True,
) -> dict[str, Any]:
    """Copy session artifacts into CreaDrug24 and register the case."""
    sid = (session_id or "").strip()
    sd = safe_session_dir(sessions_root or SESSIONS, sid)
    if not sd:
        return {"ok": False, "error": f"session not found: {sid}"}

    root = (creadrug_root or creadrug24_root()).resolve()
    if not root.is_dir():
        return {"ok": False, "error": f"CreaDrug24 root not found: {root}"}

    meta = _read_json(sd / "session_meta.json")
    s1 = _read_json(sd / "stage_01_smiles.json")
    s3 = _read_json(sd / "stage_03_index.json")
    dr = _read_json(sd / "docking_result.json")
    qc = _read_json(sd / "stage_25_qc.json")
    pocket = _read_json(sd / "stage_23_pocket_center.json") or _read_json(sd / "pocket_config.json")

    receptor_src = _first_existing(sd, ["protein/cleaned.pdb", "protein/apo_for_docking.pdb"])
    pose_pdbqt = _first_existing(sd, ["docking/output.pdbqt", "best_pose.pdbqt"])
    pose_pdb = _first_existing(sd, ["best_pose.pdb"])

    if not receptor_src:
        return {"ok": False, "error": "缺少 receptor（protein/cleaned.pdb）"}
    if not pose_pdbqt and not pose_pdb:
        return {"ok": False, "error": "缺少 docking pose（docking/output.pdbqt 或 best_pose.pdb）"}
    if not dr:
        return {"ok": False, "error": "缺少 docking_result.json（請先跑完 Module 5 Engine A）"}

    drug_id = _resolve_drug_id(sd, meta, s3)
    cid = _sanitize_case_id(case_id or _default_case_id(sid, drug_id))
    case_dir = root / "input" / "cases" / cid
    if case_dir.exists() and not overwrite:
        return {"ok": False, "error": f"case 已存在：{cid}（加 overwrite=true 覆寫）"}
    case_dir.mkdir(parents=True, exist_ok=True)

    shutil.copy2(receptor_src, case_dir / "receptor.pdb")
    pose_ext = "pdbqt"
    if pose_pdbqt:
        shutil.copy2(pose_pdbqt, case_dir / "best_pose.pdbqt")
        pose_ext = "pdbqt"
    if pose_pdb:
        shutil.copy2(pose_pdb, case_dir / "best_pose.pdb")
        if not pose_pdbqt:
            pose_ext = "pdb"

    for src_name, dst_name in (
        ("docking_result.json", "docking_result.json"),
        ("final_report.json", "final_report.json"),
        ("stage_24_vina.json", "stage_24_vina.json"),
        ("docking/vina.log", "vina.log"),
    ):
        src = safe_session_file(sd, src_name)
        if src:
            shutil.copy2(src, case_dir / dst_name)

    if pocket:
        (case_dir / "stage_23_pocket_center.json").write_text(
            json.dumps(pocket, indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8",
        )

    dr_export = dict(dr)
    dr_export["output_pdbqt"] = f"input/cases/{cid}/best_pose.{pose_ext}"
    (case_dir / "docking_result.json").write_text(
        json.dumps(dr_export, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )

    manifest = {
        "caseId": cid,
        "session_id": sid,
        "catalogDrugId": drug_id,
        "exported_at": datetime.now(timezone.utc).isoformat(),
        "source": "drugops",
        "session_dir": str(sd),
        "creadrug_case_dir": str(case_dir),
    }
    (case_dir / "case_manifest.json").write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )

    case_dir_rel = f"input/cases/{cid}"
    index_key = str(
        (s3 or {}).get("index_key")
        or (meta or {}).get("target")
        or (dr or {}).get("target")
        or ""
    ).strip().upper()
    pdb_id = str(
        (meta or {}).get("pdb_id")
        or (s3 or {}).get("pdb_id")
        or (dr or {}).get("pdb_id")
        or ""
    ).strip().upper()
    smiles = str(
        (s1 or {}).get("canonical_smiles")
        or (s1 or {}).get("input_smiles")
        or (meta or {}).get("smiles_input")
        or ""
    ).strip()

    entry = build_case_entry(
        cid,
        case_dir_rel,
        drug_id=drug_id,
        drug_name=_drug_name(meta, s3, dr),
        smiles=smiles,
        target_name=_target_label(meta, s3, dr),
        index_key=index_key,
        pdb_id=pdb_id,
        dr=dr,
        qc=qc,
        vina_log=case_dir / "vina.log",
        pose_ext=pose_ext,
    )
    reg_path = _update_registry(root, entry)

    return {
        "ok": True,
        "session_id": sid,
        "case_id": cid,
        "case_dir": str(case_dir),
        "registry": str(reg_path),
        "entry": entry,
        "creadrug_root": str(root),
    }