app/services.py

import json
import os
from functools import lru_cache

from .engine import ChemicalSearchEngine
from .utils import smiles_to_image_url
from .generation import generate_explanations_batch

# Global variables
engine = None
dataset = None
index_path = None
data_path = None


def get_data_paths():
    """Get paths for data and index files."""
    global data_path, index_path
    
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    data_path = os.path.join(base_dir, "data", "compounds.json")
    index_path = os.path.join(base_dir, "data", "compounds_index.pkl")
    
    return data_path, index_path


def data_exists():
    """Check if compounds.json exists and has data."""
    data_path, _ = get_data_paths()
    
    if not os.path.exists(data_path):
        return False
    
    try:
        with open(data_path) as f:
            data = json.load(f)
        return len(data) > 0
    except:
        return False


def index_exists():
    """Check if FAISS index exists."""
    _, index_path = get_data_paths()
    return os.path.exists(index_path)


def initialize_engine():
    """
    Centralized initialization:
    1. Check if compounds.json exists → use it
    2. If not → run ingest.py
    3. Check if FAISS index exists → load it
    4. If not → build it
    """
    global engine, dataset
    
    if engine is not None and dataset is not None:
        return engine
    
    data_path, index_path = get_data_paths()
    
    print("\n" + "="*60)
    print("[STARTUP] Initializing Chemical RAG System (Centralized)")
    print("="*60 + "\n")
    
    # Step 1: Check for compounds.json
    if not data_exists():
        print("[WARNING] compounds.json not found or empty")
        print("[INFO] Running ingestion pipeline...")
        from . import ingest_handler
        ingest_handler.run_ingestion()
    
    # Step 2: Load compounds.json
    print(f"[LOAD] Loading compounds from {data_path}...")
    try:
        with open(data_path) as f:
            dataset = json.load(f)
        print(f"[SUCCESS] Loaded {len(dataset)} compounds")
    except Exception as e:
        raise RuntimeError(f"Failed to load compounds: {e}")
    
    # Step 3: Initialize engine
    smiles_list = [d["smiles"] for d in dataset]
    engine = ChemicalSearchEngine(bit_size=2048)
    
    # Step 4: Check for FAISS index
    if index_exists():
        print(f"[LOAD] FAISS index found at {index_path}")
        if engine.load_index(index_path):
            print("[SUCCESS] FAISS-IVF index loaded successfully")
            print(f"[SUCCESS] Generation layer enabled with Llama-3.1-8B")
            print("="*60 + "\n")
            return engine
    
    # Step 5: Build and save FAISS index
    print("[BUILD] Building FAISS-IVF index (this may take a few minutes)...")
    engine.add_compounds(smiles_list, metadata_list=dataset)
    engine.save_index(index_path)
    
    print(f"[SUCCESS] Generation layer enabled with Llama-3.1-8B")
    print(f"[SUCCESS] FAISS index saved to {index_path}")
    print("="*60 + "\n")
    return engine


def _search_internal(smiles: str, top_k: int, base_url: str = None, include_explanation: bool = True):
    """
    Internal search using FAISS-IVF with optional LLM explanation.
    
    Args:
        smiles: Query SMILES string
        top_k: Number of results (1-100)
        base_url: Base URL for image URLs (e.g. https://example.com)
        include_explanation: Whether to generate LLM explanations
    
    Returns:
        List of enriched results with metadata and optional explanations
    """
    global engine, dataset
    
    if engine is None:
        raise RuntimeError("Engine not initialized. Call initialize_engine() first.")
    
    if dataset is None:
        raise RuntimeError("Dataset not loaded. Call initialize_engine() first.")
    
    # FAISS-IVF search
    results = engine.search(smiles, top_k)
    
    if not results:
        return []  # Empty results is valid - just no similar compounds found
    
    enriched = []
    for r in results:
        cid = r["metadata"].get("cid")
        enriched.append({
            "smiles": r["smiles"],
            "similarity_score": r["similarity_score"],
            "image": smiles_to_image_url(r["smiles"], base_url=base_url),
            "cid": str(cid) if cid is not None else None,  # Convert to string for schema
            "name": r["metadata"].get("name"),
            "mw": r["metadata"].get("mw"),
            "explanation": None
        })
    
    # Generate LLM explanations if requested
    if include_explanation:
        enriched = generate_explanations_batch(smiles, enriched)
    
    return enriched


@lru_cache(maxsize=1000)
def cached_search(smiles: str, top_k: int, explain: bool = True):
    """
    Cached search results using FAISS-IVF with optional explanations.
    """
    results = _search_internal(smiles, top_k, include_explanation=explain)
    return tuple([tuple(sorted(r.items())) for r in results])


def get_search_results(smiles: str, top_k: int = 3, explain: bool = True, base_url: str = None):
    """
    Main search function: FAISS-IVF retrieval with optional LLM generation.
    
    Args:
        smiles: Query SMILES string
        top_k: Number of results to return (default 3, max 100)
        explain: Whether to generate LLM explanations (default True)
        base_url: Base URL for image URLs (e.g. https://example.com)
    
    Returns:
        Tuple of (results_list, query_smiles)
    """
    results = _search_internal(smiles, top_k, base_url=base_url, include_explanation=explain)
    return results, smiles


def get_search_results_retrieval_only(smiles: str, top_k: int = 3, base_url: str = None):
    """
    Fast retrieval-only search (no LLM generation).
    
    Args:
        smiles: Query SMILES string
        top_k: Number of results to return (default 3, max 100)
        base_url: Base URL for image URLs (e.g. https://example.com)
    
    Returns:
        Tuple of (results_list, query_smiles)
    """
    results = _search_internal(smiles, top_k, base_url=base_url, include_explanation=False)
    return results, smiles


def get_system_stats():
    """Get system statistics and status."""
    global dataset, engine
    
    if engine is None or dataset is None:
        return {
            "status": "uninitialized",
            "compounds": 0,
            "index_type": "FAISS-BinaryFlat",
            "index_exists": False
        }
    
    return {
        "status": "ready",
        "compounds": len(dataset),
        "index_type": "FAISS-BinaryFlat (Binary Flat Index)",
        "index_built": engine.index_built,
        "total_indexed": engine.total_compounds,
        "fingerprint_bits": engine.bit_size,
        "similarity_metric": "Tanimoto (exact)",
        "generation_model": "Llama-3.1-8B-Instruct",
        "generation_enabled": True,
        "endpoints": [
            "/search/retrieval-only - Fast retrieval only",
            "/search/full-rag - Retrieval + LLM explanation",
            "/stats - System statistics",
            "/health - Health check"
        ]
    }