#!/usr/bin/env python3
"""Encode an amino-acid sequence to a ProRiboGen-compatible VESM3B H5."""
from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

from common import PKG_ROOT, device, load_env, resolve, vesm_paths


def write_fasta(path: Path, p_id: str, sequence: str) -> None:
    seq = re.sub(r"\s+", "", sequence).upper()
    if not seq or not re.fullmatch(r"[ACDEFGHIKLMNPQRSTVWY*BXZ]+", seq):
        raise ValueError("Protein sequence is empty or contains invalid characters")
    path.write_text(f">{p_id}\n{seq}\n", encoding="utf-8")


def encode_protein_to_h5(
    protein_seq: str,
    *,
    p_id: str = "QUERY",
    output_h5: Path | None = None,
    fp16: bool = True,
) -> Path:
    load_env()
    base, weights = vesm_paths()
    if not base.is_dir():
        raise FileNotFoundError(f"VESM 基座不存在: {base}")
    if not weights.is_file():
        raise FileNotFoundError(f"VESM 权重不存在: {weights}")

    work = resolve(os.environ.get("WORKSPACE", "workspace"))
    work.mkdir(parents=True, exist_ok=True)
    out = output_h5 or (work / f"{p_id}_vesm3b.h5")

    with tempfile.TemporaryDirectory(dir=work) as tmp:
        fa = Path(tmp) / f"{p_id}.fasta"
        write_fasta(fa, p_id, protein_seq)
        script = PKG_ROOT / "protein_encoder" / "build_vesm3b_protein_embeddings_h5.py"
        cmd = [
            sys.executable,
            str(script),
            "--base-model-dir",
            str(base),
            "--vesm-weights",
            str(weights),
            "--fasta",
            str(fa),
            "--output",
            str(out),
            "--device",
            device(),
        ]
        if fp16:
            cmd.append("--fp16")
        subprocess.run(cmd, check=True)
    return out


def main() -> None:
    ap = argparse.ArgumentParser(description="Protein AA sequence → VESM3B H5")
    ap.add_argument("--protein", required=True, help="AA sequence, or @path/to.fasta")
    ap.add_argument("--p-id", default="QUERY")
    ap.add_argument("--output", type=Path, default=None)
    ap.add_argument("--no-fp16", action="store_true")
    args = ap.parse_args()

    load_env()
    protein = args.protein
    if protein.startswith("@"):
        text = Path(protein[1:]).read_text(encoding="utf-8")
        lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
        if lines and lines[0].startswith(">"):
            protein = "".join(lines[1:])
            if args.p_id == "QUERY":
                args.p_id = lines[0][1:].split()[0]
        else:
            protein = "".join(lines)

    out = encode_protein_to_h5(
        protein, p_id=args.p_id, output_h5=args.output, fp16=not args.no_fp16
    )
    print(f"OK -> {out}")


if __name__ == "__main__":
    main()