app/main.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.staticfiles import StaticFiles
from fastapi.concurrency import run_in_threadpool
import os

from .schemas import SearchRequest, SearchResponse, CompoundResult
from .services import (
    initialize_engine, 
    get_search_results, 
    get_search_results_retrieval_only,
    get_system_stats
)

app = FastAPI(
    title="Chemical RAG System v2.1",
    version="2.1.0",
    description="FAISS-IVF Retrieval-Augmented Generation for 1M+ chemical compounds"
)

# Mount static files
if os.path.exists("app/static"):
    app.mount("/static", StaticFiles(directory="app/static"), name="static")


@app.on_event("startup")
async def startup_event():
    """Initialize the engine on startup with centralized logic."""
    try:
        initialize_engine()
        print("[SUCCESS] API startup successful (v2.1.0)")
    except Exception as e:
        import traceback
        print(f"[ERROR] Startup failed: {str(e)}")
        traceback.print_exc()


@app.get("/")
async def root():
    """Health check endpoint."""
    return {
        "status": "running",
        "service": "Chemical RAG System with FAISS-IVF",
        "version": "2.1.0",
        "endpoints": {
            "/search/retrieval-only": "Fast retrieval using FAISS-IVF (no LLM)",
            "/search/full-rag": "Full RAG pipeline with LLM explanation",
            "/stats": "System statistics",
            "/health": "Health check"
        }
    }


@app.get("/health")
async def health():
    """Health check with detailed status."""
    stats = get_system_stats()
    return {
        "status": "healthy",
        "service": "Chemical RAG System",
        "version": "2.1.0",
        "system": stats,
        "features": [
            "FAISS-IVF Indexing (1M+ compound support)",
            "Fast retrieval (<100ms)",
            "LLM Explanations (Optional)",
            "Chemical accuracy preserved"
        ]
    }


@app.post("/search/retrieval-only", response_model=SearchResponse)
async def search_retrieval_only(payload: SearchRequest, http_request: Request):
    """
    ⚡ FAST RETRIEVAL ENDPOINT (No LLM generation)
    
    Uses FAISS-IVF for ultra-fast chemical similarity search.
    
    Performance:
    - 1M compounds: <100ms
    - No LLM overhead
    - Chemical accuracy maintained
    
    Response includes:
    - SMILES and similarity scores
    - Compound metadata (name, CID, MW)
    - No explanations
    """
    
    # Validate SMILES
    if not payload.smiles or len(payload.smiles.strip()) == 0:
        raise HTTPException(status_code=400, detail="SMILES string cannot be empty")
    
    if payload.top_k < 1 or payload.top_k > 100:
        raise HTTPException(status_code=400, detail="top_k must be between 1 and 100")

    try:
        # Get base URL from request
        base_url = str(http_request.base_url).rstrip('/')
        
        # Run retrieval-only search (no generation)
        results, query_smiles = await run_in_threadpool(
            get_search_results_retrieval_only,
            payload.smiles.strip(),
            payload.top_k,
            base_url
        )

        # Convert to response model (empty results is OK - just no matches found)
        compound_results = [
            CompoundResult(
                smiles=r["smiles"],
                similarity_score=r["similarity_score"],
                image=r.get("image"),
                explanation=None,  # Retrieval-only mode
                cid=r.get("cid"),
                name=r.get("name")
            )
            for r in results
        ]

        return SearchResponse(
            results=compound_results,
            query_smiles=query_smiles,
            total_results=len(compound_results)
        )

    except HTTPException:
        raise
    except Exception as e:
        # Check if it's a SMILES validation issue
        error_msg = str(e).lower()
        if "smiles" in error_msg or "invalid" in error_msg:
            raise HTTPException(status_code=400, detail=f"Invalid SMILES string: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/search/full-rag", response_model=SearchResponse)
async def search_full_rag(payload: SearchRequest, http_request: Request):
    """
    🤖 FULL RAG ENDPOINT (Retrieval + LLM Explanation)
    
    Combines FAISS-IVF retrieval with Llama-3.1-8B explanations.
    
    Performance:
    - 1M compounds: <500ms (FAISS + LLM)
    - Full RAG pipeline
    - Chemical explanations included
    
    Response includes:
    - SMILES and similarity scores
    - Compound metadata (name, CID, MW)
    - LLM-generated explanations of why compounds are similar
    """
    
    # Validate SMILES
    if not payload.smiles or len(payload.smiles.strip()) == 0:
        raise HTTPException(status_code=400, detail="SMILES string cannot be empty")
    
    if payload.top_k < 1 or payload.top_k > 100:
        raise HTTPException(status_code=400, detail="top_k must be between 1 and 100")

    try:
        # Get base URL from request
        base_url = str(http_request.base_url).rstrip('/')
        
        # Run full RAG search with explanations
        results, query_smiles = await run_in_threadpool(
            get_search_results,
            payload.smiles.strip(),
            payload.top_k,
            payload.explain,  # Use the explain parameter
            base_url
        )

        # Convert to response model (empty results is OK - just no matches found)
        compound_results = [
            CompoundResult(
                smiles=r["smiles"],
                similarity_score=r["similarity_score"],
                image=r.get("image"),
                explanation=r.get("explanation"),  # LLM explanation included
                cid=r.get("cid"),
                name=r.get("name")
            )
            for r in results
        ]

        return SearchResponse(
            results=compound_results,
            query_smiles=query_smiles,
            total_results=len(compound_results)
        )

    except HTTPException:
        raise
    except Exception as e:
        # Check if it's a SMILES validation issue
        error_msg = str(e).lower()
        if "smiles" in error_msg or "invalid" in error_msg:
            raise HTTPException(status_code=400, detail=f"Invalid SMILES string: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/stats")
async def stats():
    """Get system statistics including FAISS-IVF index info."""
    system_stats = get_system_stats()
    return system_stats