from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import sys
import os
import pandas as pd

# Ensure src modules can be imported
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))

try:
    from inference import ABSAPredictor
    from rag_engine import ReviewRAGEngine
except ImportError as e:
    print(f"Error importing core modules: {e}")
    ABSAPredictor = None
    ReviewRAGEngine = None

app = FastAPI(
    title="Aspect-Based Review Intelligence API",
    description="Backend for ABSA predictions and RAG-based Review Q&A",
    version="1.0"
)

# Initialize engines lazily on startup
absa_engine = None
rag_engine = None

@app.on_event("startup")
def load_engines():
    global absa_engine, rag_engine
    print("Loading models into memory...")
    if ABSAPredictor:
        absa_engine = ABSAPredictor(model_path='models/absa-roberta-final')
    if ReviewRAGEngine:
        rag_engine = ReviewRAGEngine(vectorstore_dir='vectorstore')

class ReviewRequest(BaseModel):
    text: str
    
class QuestionRequest(BaseModel):
    question: str
    top_k: int = 15

@app.get("/")
def health_check():
    return {
        "status": "online",
        "absa_loaded": absa_engine is not None and absa_engine.is_loaded,
        "rag_loaded": rag_engine is not None and rag_engine.index is not None
    }

@app.post("/predict")
def predict_aspects(req: ReviewRequest):
    if not absa_engine:
        raise HTTPException(status_code=503, detail="ABSA engine failed to initialize")
        
    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Review text cannot be empty")
        
    results = absa_engine.predict(req.text)
    return {
        "review_text": req.text,
        "aspect_sentiments": results
    }

@app.post("/ask")
def ask_question(req: QuestionRequest):
    if not rag_engine:
        raise HTTPException(status_code=503, detail="RAG engine failed to initialize")
        
    if not req.question.strip():
        raise HTTPException(status_code=400, detail="Question cannot be empty")
        
    answer = rag_engine.answer_question(req.question, top_k=req.top_k)
    return {
        "question": req.question,
        "answer": answer
    }
    
@app.get("/stats")
def get_aggregate_stats():
    """Returns basic aggregate statistics from the FAISS metadata database for dashboard plots"""
    if not rag_engine or rag_engine.metadata_df is None:
        return {"error": "Metadata database not found in vectorstore."}
        
    df = rag_engine.metadata_df
    
    aspect_counts = {
        'food': {'positive': 0, 'negative': 0, 'neutral': 0, 'conflict': 0},
        'service': {'positive': 0, 'negative': 0, 'neutral': 0, 'conflict': 0},
        'ambiance': {'positive': 0, 'negative': 0, 'neutral': 0, 'conflict': 0},
        'price': {'positive': 0, 'negative': 0, 'neutral': 0, 'conflict': 0},
        'anecdotes/miscellaneous': {'positive': 0, 'negative': 0, 'neutral': 0, 'conflict': 0}
    }
    
    import json
    for aspects_str in df.get('predicted_aspects', []):
        if pd.isna(aspects_str) or not aspects_str:
            continue
        try:
            aspect_dict = json.loads(aspects_str)
            for aspect, data in aspect_dict.items():
                if aspect in aspect_counts and data['sentiment'] in aspect_counts[aspect]:
                    aspect_counts[aspect][data['sentiment']] += 1
        except:
             continue
             
    return {
        "total_reviews": len(df),
        "aspect_sentiment_counts": aspect_counts
    }

if __name__ == "__main__":
    import uvicorn
    # Typically run with: uvicorn api.main:app --reload
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)