backend/app.py

import os
import sys
import json
import socket
import traceback
from pathlib import Path
from typing import Dict, Any

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

from rdkit import Chem
from rdkit.Chem import Descriptors


# ============================================================
# Path 설정
# ============================================================

BASE_DIR = Path(__file__).resolve().parents[1]

MODELS_DIR = BASE_DIR / "models"
PBPK_ENGINE_DIR = BASE_DIR / "pbpk_engine"
RESULTS_DIR = PBPK_ENGINE_DIR / "results"

# React 빌드 경로 (Dockerfile에서 빌드됨)
FRONTEND_BUILD_DIR = BASE_DIR / "frontend" / "dist"

ADMET_FT_DIR = MODELS_DIR / "admet_ft"
CLEARANCE_FT_DIR = MODELS_DIR / "clearance_ft"

ADMET_MODEL_PATH = ADMET_FT_DIR / "results" / "final_model"

sys.path.insert(0, str(MODELS_DIR))
sys.path.insert(0, str(ADMET_FT_DIR))
sys.path.insert(0, str(CLEARANCE_FT_DIR))

from admet_ft.inference import load_admet_model, admet_prediction
from admet_ft._modules.rdkit import load_rdkit_description
from clearance_ft.main_prediction import load_clearance_models, run_clearance


# ============================================================
# FastAPI
# ============================================================

app = FastAPI(title="PharmAI Unified API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# /results/... → PBPK 결과 파일 (CSV, PNG) 서빙
RESULTS_DIR.mkdir(exist_ok=True)
app.mount("/results", StaticFiles(directory=str(RESULTS_DIR)), name="results")

# /assets/... → React 빌드 정적 파일 서빙 (JS, CSS)
if FRONTEND_BUILD_DIR.exists():
    app.mount("/assets", StaticFiles(directory=str(FRONTEND_BUILD_DIR / "assets")), name="assets")


# ============================================================
# Global models
# ============================================================

ADMET_MODEL = None
CLEARANCE_MODELS = None

R_SERVER_HOST = "127.0.0.1"
R_SERVER_PORT = 7000


# ============================================================
# Utilities
# ============================================================

def calculate_molecular_weight(smiles: str) -> float:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Invalid SMILES: {smiles}")
    return float(Descriptors.MolWt(mol))


def run_clearance_safe(smiles: str) -> float:
    if CLEARANCE_MODELS is None:
        raise RuntimeError("Clearance model is not loaded.")

    original_dir = os.getcwd()

    try:
        os.chdir(str(CLEARANCE_FT_DIR))

        result = run_clearance(
            smiles,
            preloaded_models=CLEARANCE_MODELS,
        )

        return float(result["Clint"].iloc[0])

    finally:
        os.chdir(original_dir)


def send_to_r_server(drug_data: Dict[str, Any]) -> Dict[str, Any]:
    request = {"data": drug_data}

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(120)

    try:
        sock.connect((R_SERVER_HOST, R_SERVER_PORT))
        sock.sendall((json.dumps(request) + "\n").encode("utf-8"))

        total_data = b""

        while True:
            chunk = sock.recv(8192)
            if not chunk:
                break

            total_data += chunk

            if b"\n" in total_data:
                break

        if not total_data:
            raise RuntimeError("No response from R PBPK server.")

        return json.loads(total_data.decode("utf-8").strip())

    finally:
        sock.close()


# ============================================================
# Startup
# ============================================================

@app.on_event("startup")
async def startup_event():
    global ADMET_MODEL, CLEARANCE_MODELS

    print("=" * 60)
    print("Starting PharmAI Unified API")
    print("=" * 60)

    try:
        print("[1/2] Loading ADMET property model...")
        print(f"ADMET_MODEL_PATH = {ADMET_MODEL_PATH}")

        if not ADMET_MODEL_PATH.exists():
            raise FileNotFoundError(f"ADMET model path not found: {ADMET_MODEL_PATH}")

        ADMET_MODEL, _ = load_admet_model(str(ADMET_MODEL_PATH))
        print("✅ ADMET property model loaded")

    except Exception as e:
        print(f"❌ ADMET property model load failed: {e}")
        traceback.print_exc()
        ADMET_MODEL = None

    try:
        print("[2/2] Loading clearance models...")

        original_dir = os.getcwd()

        try:
            os.chdir(str(CLEARANCE_FT_DIR))
            CLEARANCE_MODELS = load_clearance_models()
        finally:
            os.chdir(original_dir)

        print("✅ Clearance models loaded")

    except Exception as e:
        print(f"❌ Clearance model load failed: {e}")
        traceback.print_exc()
        CLEARANCE_MODELS = None

    print("=" * 60)
    print("PharmAI Unified API ready")
    print(f"ADMET loaded: {ADMET_MODEL is not None}")
    print(f"Clearance loaded: {CLEARANCE_MODELS is not None}")
    print(f"R server: {R_SERVER_HOST}:{R_SERVER_PORT}")
    print("=" * 60)


# ============================================================
# Endpoints
# ============================================================

@app.get("/api/health")
async def health():
    return {
        "status": "ok",
        "message": "PharmAI Unified API is running.",
    }


@app.get("/api/status")
async def status():
    return {
        "status": "running",
        "admet_loaded": ADMET_MODEL is not None,
        "clearance_loaded": CLEARANCE_MODELS is not None,
        "r_server_host": R_SERVER_HOST,
        "r_server_port": R_SERVER_PORT,
    }


@app.post("/api/predict/admet")
async def predict_admet(data: dict):
    smiles = data.get("smiles", "").strip()

    if not smiles:
        raise HTTPException(status_code=400, detail="SMILES is required.")

    if ADMET_MODEL is None:
        raise HTTPException(status_code=503, detail="ADMET model is not loaded.")

    try:
        result = admet_prediction(
            smiles,
            model_path=str(ADMET_MODEL_PATH),
            trainer=ADMET_MODEL,
        )

        rdkit_desc = load_rdkit_description(smiles)
        result["molecular_weight"] = float(rdkit_desc["MolWt"])

        # frontend가 바로 pKa, logP 등을 받도록 그대로 반환
        return result

    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"ADMET prediction failed: {e}")


@app.post("/api/predict/clearance")
async def predict_clearance(data: dict):
    smiles = data.get("smiles", "").strip()

    if not smiles:
        raise HTTPException(status_code=400, detail="SMILES is required.")

    try:
        clint = run_clearance_safe(smiles)

        return {
            "SMILES": smiles,
            "Clint": clint,
        }

    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Clearance prediction failed: {e}")


@app.post("/api/pbpk/simulate")
async def simulate_pbpk(data: dict):
    smiles = data.get("smiles", "").strip()
    drug_name = data.get("drug_name", "Unknown_Compound")

    use_admet = data.get("use_admet", True)
    use_clearance = data.get("use_clearance", True)

    if not smiles:
        raise HTTPException(status_code=400, detail="SMILES is required.")

    try:
        drug_data = {
            "Drug.Name": drug_name,
            "Molecular.weight": calculate_molecular_weight(smiles),
            "pKa": None,
            "logP": None,
            "fu_in_vitro": None,
            "permeability": None,
            "solubility": None,
            "Clint": None,
        }

        if use_admet:
            if ADMET_MODEL is None:
                raise HTTPException(status_code=503, detail="ADMET model is not loaded.")

            admet_result = admet_prediction(
                smiles,
                model_path=str(ADMET_MODEL_PATH),
                trainer=ADMET_MODEL,
            )

            drug_data["pKa"] = admet_result.get("pKa", 7.4)
            drug_data["logP"] = admet_result.get("logP", 2.0)
            drug_data["fu_in_vitro"] = admet_result.get("fu_in_vitro", 0.5)
            drug_data["permeability"] = admet_result.get("permeability", 10.0)
            drug_data["solubility"] = admet_result.get("solubility", 100.0)

        else:
            db_features = data.get("db_features", {})
            drug_data["pKa"] = db_features.get("pKa", 7.4)
            drug_data["logP"] = db_features.get("logP", 2.0)
            drug_data["fu_in_vitro"] = db_features.get("fu_in_vitro", 0.5)
            drug_data["permeability"] = db_features.get("permeability", 10.0)
            drug_data["solubility"] = db_features.get("solubility", 100.0)

        if use_clearance:
            drug_data["Clint"] = run_clearance_safe(smiles)
        else:
            db_features = data.get("db_features", {})
            drug_data["Clint"] = db_features.get("Clint", 980.0)

        pbpk_result = send_to_r_server(drug_data)

        result_paths = {
            "nca_csv": f"results/{drug_name}_PBPK_Results/{drug_name}_NCA_results.csv",
            "conc_csv": f"results/{drug_name}_PBPK_Results/{drug_name}_conc_results.csv",
            "pbpk_plot": f"results/{drug_name}_PBPK_Results/{drug_name}_PBPK_figure.png",
            "acat_plot": f"results/{drug_name}_PBPK_Results/{drug_name}_ACAT_figure.png",
        }

        return {
            "status": "success",
            "message": "PBPK simulation completed successfully",
            "drug_name": drug_name,
            "drug_data": drug_data,
            "result_paths": result_paths,
            "pbpk_result": pbpk_result,
        }

    except HTTPException:
        raise

    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"PBPK simulation failed: {e}")


# 임시 endpoint: 프론트 Index.tsx가 /api/predict를 호출하므로 서버 에러 방지용
# CYP 모델 붙이기 전까지는 placeholder
@app.post("/api/predict")
async def predict_classification_placeholder(data: dict):
    smiles = data.get("smiles", "").strip()

    if not smiles:
        raise HTTPException(status_code=400, detail="SMILES is required.")

    return {
        "smiles": smiles,
        "predictions": {},
        "meta": {
            "message": "CYP/ADMET classification model is not connected yet."
        },
    }


# ============================================================
# React SPA 라우팅: /api, /results 이외 모든 경로 → index.html
# (이 mount는 반드시 API 엔드포인트 정의 후 마지막에 위치해야 함)
# ============================================================

@app.get("/{full_path:path}")
async def serve_react_app(full_path: str, request: Request):
    """React SPA: API가 아닌 모든 경로는 index.html로 처리"""
    index_file = FRONTEND_BUILD_DIR / "index.html"
    if FRONTEND_BUILD_DIR.exists() and index_file.exists():
        return FileResponse(str(index_file))
    return {"message": "PharmAI API is running. Frontend build not found."}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=7860,  # HuggingFace Spaces 기본 포트
        reload=False,
    )