from pathlib import Path
from dataclasses import dataclass

import requests


@dataclass
class MSAResult:
    """Multiple Sequence Alignment 결과."""
    query_id: str
    a3m_content: str  # a3m 포맷 MSA
    depth: int         # MSA 깊이 (동원된 서열 수)


def build_msa_remote(
    query_sequence: str,
    query_id: str = "target",
    api_base: str = "https://api.colabfold.com",
    max_depth: int = 5000,
) -> MSAResult:
    """ColabFold MSA 서버(MMseqs2 백엔드)로 원격 MSA 구축.
    
    Boltz 공식 문서는 ColabFold 또는 BioLM API를 권장 [4].
    """
    # 실제 API 호출 (예시: ColabFold MMseqs2 API)
    payload = {"query": f">{query_id}\n{query_sequence}", "mode": "env"}
    resp = requests.post(f"{api_base}/ticket/msa", data=payload, timeout=600)
    resp.raise_for_status()
    ticket = resp.json()["id"]
    
    # 폴링해서 완료 대기
    import time
    while True:
        status = requests.get(f"{api_base}/ticket/{ticket}", timeout=30).json()
        if status["status"] == "COMPLETE":
            break
        elif status["status"] == "ERROR":
            raise RuntimeError(f"MSA 구축 실패: {status}")
        time.sleep(5)
    
    # a3m 결과 다운로드
    a3m_resp = requests.get(f"{api_base}/result/download/{ticket}", timeout=60)
    a3m_resp.raise_for_status()
    a3m_content = a3m_resp.text
    depth = a3m_content.count("\n>")
    
    return MSAResult(query_id=query_id, a3m_content=a3m_content, depth=depth)


def save_a3m(msa: MSAResult, output_dir: Path) -> Path:
    """a3m 파일 저장."""
    output_dir.mkdir(parents=True, exist_ok=True)
    path = output_dir / f"{msa.query_id}.a3m"
    path.write_text(msa.a3m_content)
    return path

from rdkit import Chem
from rdkit.Chem import AllChem, Descriptors, Lipinski
import pandas as pd


def canonicalize_smiles(smiles: str) -> str | None:
    """SMILES canonical form 반환. 파싱 실패 시 None."""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    return Chem.MolToSmiles(mol, canonical=True)


def lipinski_filter(smiles: str) -> dict:
    """Lipinski's Rule of Five 통과 여부 + 물리화학 특성."""
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return {"valid": False}
    props = {
        "MW": Descriptors.MolWt(mol),
        "LogP": Descriptors.MolLogP(mol),
        "HBA": Lipinski.NumHAcceptors(mol),
        "HBD": Lipinski.NumHDonors(mol),
        "RotB": Lipinski.NumRotatableBonds(mol),
    }
    violations = sum([
        props["MW"] > 500,
        props["LogP"] > 5,
        props["HBA"] > 10,
        props["HBD"] > 5,
    ])
    props["valid"] = True
    props["lipinski_violations"] = violations
    props["lipinski_pass"] = violations <= 1
    return props


def prepare_ligand_library(smiles_list: list[str]) -> pd.DataFrame:
    """리간드 라이브러리 정규화 + 물리화학 필터."""
    records = []
    for i, raw in enumerate(smiles_list):
        canon = canonicalize_smiles(raw)
        if canon is None:
            continue
        props = lipinski_filter(canon)
        records.append({
            "ligand_id": f"L{i:04d}",
            "smiles_raw": raw,
            "smiles_canonical": canon,
            **props,
        })
    df = pd.DataFrame(records)
    return df[df["lipinski_pass"]].reset_index(drop=True)

import yaml


def build_boltz_config(
    protein_sequence: str,
    protein_msa_path: Path,
    ligand_smiles: str,
    ligand_id: str,
    output_dir: Path,
) -> Path:
    """Boltz YAML config 생성. 하나의 protein-ligand 페어 당 하나."""
    config = {
        "version": 1,
        "sequences": [
            {
                "protein": {
                    "id": "A",
                    "sequence": protein_sequence,
                    "msa": str(protein_msa_path),
                }
            },
            {
                "ligand": {
                    "id": "L",
                    "smiles": ligand_smiles,
                }
            },
        ],
        "constraints": [],
        "properties": [
            {"affinity": {"binder": "L"}}  # 결합 친화도 예측 켜기 (Boltz-2 신규)
        ],
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    yaml_path = output_dir / f"{ligand_id}.yaml"
    with open(yaml_path, "w") as f:
        yaml.safe_dump(config, f, sort_keys=False)
    return yaml_path

import subprocess
import json


@dataclass
class BoltzResult:
    """Boltz-2 예측 결과."""
    ligand_id: str
    pdb_path: Path            # 예측된 3D 복합체 구조
    plddt_mean: float         # 구조 신뢰도 (0~100)
    iptm: float               # 인터페이스 pTM (0~1)
    affinity_log_ic50: float  # log10(IC50 mol/L) 예측치 (음수일수록 강력)
    inference_time_sec: float


def run_boltz_prediction(
    yaml_path: Path,
    output_dir: Path,
    device: str = "cuda",
    use_msa_server: bool = False,
) -> BoltzResult:
    """Boltz-2 CLI 실행 (실전에서는 Python API도 가능)."""
    import time
    start = time.time()
    
    cmd = [
        "boltz", "predict",
        str(yaml_path),
        "--out_dir", str(output_dir),
        "--devices", "1",
        "--accelerator", "gpu" if device == "cuda" else "cpu",
    ]
    if use_msa_server:
        cmd.append("--use_msa_server")
    
    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    elapsed = time.time() - start
    
    # Boltz는 output_dir/{name}/{name}_model_0.pdb + confidence.json 생성
    name = yaml_path.stem
    pdb_path = output_dir / name / f"{name}_model_0.pdb"
    conf_path = output_dir / name / f"confidence_{name}_model_0.json"
    
    with open(conf_path) as f:
        conf = json.load(f)
    
    # affinity는 별도 JSON에서 (Boltz-2 신규 필드)
    affinity_path = output_dir / name / f"affinity_{name}.json"
    affinity_data = json.loads(affinity_path.read_text()) if affinity_path.exists() else {}
    
    return BoltzResult(
        ligand_id=name,
        pdb_path=pdb_path,
        plddt_mean=float(conf.get("complex_plddt", 0.0)),
        iptm=float(conf.get("iptm", 0.0)),
        affinity_log_ic50=float(affinity_data.get("affinity_pred_value", 0.0)),
        inference_time_sec=elapsed,
    )

from tqdm import tqdm


def batch_screen(
    protein_sequence: str,
    protein_msa: MSAResult,
    ligand_library: pd.DataFrame,
    work_dir: Path,
    device: str = "cuda",
) -> pd.DataFrame:
    """리간드 라이브러리 전체 스크리닝."""
    msa_path = save_a3m(protein_msa, work_dir / "msa")
    
    results = []
    failures = []
    for _, row in tqdm(ligand_library.iterrows(), total=len(ligand_library), desc="Screening"):
        try:
            yaml_path = build_boltz_config(
                protein_sequence=protein_sequence,
                protein_msa_path=msa_path,
                ligand_smiles=row["smiles_canonical"],
                ligand_id=row["ligand_id"],
                output_dir=work_dir / "configs",
            )
            boltz_result = run_boltz_prediction(
                yaml_path=yaml_path,
                output_dir=work_dir / "predictions",
                device=device,
                use_msa_server=False,
            )
            results.append({
                "ligand_id": boltz_result.ligand_id,
                "smiles": row["smiles_canonical"],
                "affinity_log_ic50": boltz_result.affinity_log_ic50,
                "plddt": boltz_result.plddt_mean,
                "iptm": boltz_result.iptm,
                "pdb_path": str(boltz_result.pdb_path),
                "inference_sec": boltz_result.inference_time_sec,
                "MW": row.get("MW"),
                "LogP": row.get("LogP"),
            })
        except Exception as e:
            failures.append({"ligand_id": row["ligand_id"], "error": str(e)})
    
    df_results = pd.DataFrame(results).sort_values("affinity_log_ic50")
    df_failures = pd.DataFrame(failures)
    df_failures.to_csv(work_dir / "failures.csv", index=False)
    return df_results

def render_pymol(
    pdb_path: Path,
    output_png: Path,
    ray_trace: bool = True,
) -> None:
    """PyMOL 헤드리스 렌더링."""
    script = f"""
load {pdb_path}, complex
hide everything
show cartoon, complex and polymer
show sticks, complex and organic
color grey70, complex and polymer
color yellow, complex and organic
bg_color white
zoom complex and organic, 5
{"ray 1200, 900" if ray_trace else ""}
png {output_png}, dpi=150
quit
"""
    script_path = output_png.parent / "render.pml"
    script_path.write_text(script)
    subprocess.run(["pymol", "-cq", str(script_path)], check=True)


def render_top_k(results_df: pd.DataFrame, output_dir: Path, k: int = 20) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    for _, row in results_df.head(k).iterrows():
        render_pymol(
            pdb_path=Path(row["pdb_path"]),
            output_png=output_dir / f"{row['ligand_id']}.png",
        )

def full_screening_pipeline(
    protein_sequence: str,
    protein_id: str,
    smiles_library: list[str],
    work_dir: Path,
    device: str = "cuda",
    top_k: int = 20,
) -> pd.DataFrame:
    """FASTA + SMILES 라이브러리 → 랭킹된 결과 DataFrame."""
    print(f"[1/6] MSA 구축 원격 요청 (query length={len(protein_sequence)})")
    msa = build_msa_remote(protein_sequence, query_id=protein_id)
    print(f"      MSA depth={msa.depth}")
    
    print(f"[2/6] 리간드 라이브러리 정규화 (raw={len(smiles_library)})")
    library = prepare_ligand_library(smiles_library)
    print(f"      Lipinski 통과={len(library)}")
    
    print(f"[3/6] 배치 스크리닝 시작")
    results = batch_screen(protein_sequence, msa, library, work_dir, device)
    results.to_csv(work_dir / "screening_results.csv", index=False)
    
    print(f"[4/6] Top-{top_k} PyMOL 렌더링")
    render_top_k(results, work_dir / "top_k_renders", k=top_k)
    
    print(f"[5/6] 완료: work_dir={work_dir}")
    print(f"[6/6] 상위 5개 리간드 요약:")
    print(results.head(5)[["ligand_id", "affinity_log_ic50", "plddt", "iptm"]].to_string(index=False))
    return results

from chai_lab.chai1 import run_inference


def chai_reverification(
    protein_sequence: str,
    ligand_smiles: str,
    output_dir: Path,
    device: str = "cuda",
) -> dict:
    """Chai-1로 동일 페어 재예측. Boltz 결과와 앙상블 스코어링."""
    fasta_str = f">protein|name=A\n{protein_sequence}\n>ligand|name=L\n{ligand_smiles}\n"
    fasta_path = output_dir / "chai_input.fasta"
    fasta_path.write_text(fasta_str)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    result = run_inference(
        fasta_file=fasta_path,
        output_dir=output_dir,
        num_trunk_recycles=3,
        num_diffn_timesteps=200,
        seed=42,
        device=device,
        use_esm_embeddings=True,
    )
    return {
        "pdb_path": result[0],
        "iptm": float(result[0].iptm),  # 예시 필드, 실제 API 참조
    }


def ensemble_ranking(
    boltz_score: float,
    chai_score: float,
    weight: float = 0.5,
) -> float:
    """두 모델의 스코어 가중 평균 (낮을수록 강력)."""
    return weight * boltz_score + (1 - weight) * chai_score