from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
from fastapi.responses import FileResponse
from pathlib import Path
import os
import shutil

# RAG imports
from rag.qa import embed_query, answer_question_with_llm
from rag.report import build_report_from_query
from rag.vectorstore import FaissStore
from rag.render import render_docx
from rag import config
from rag.ingest import ingest_folder_incremental

app = FastAPI()

# Load existing FAISS store
store = None
if os.path.exists(config.INDEX_DIR):
    store = FaissStore.load(config.INDEX_DIR)

# Q&A models
class QuestionRequest(BaseModel):
    question: str
    k: int = 15

class AnswerResponse(BaseModel):
    question: str
    answer: str
    sources_used: int

# Report models
class ReportRequest(BaseModel):
    query: str
    filename: str = "report.docx"
    k: int = 100

class ReportResponse(BaseModel):
    message: str
    filename: str
    download_url: str

UPLOAD_DIR = Path("data")
UPLOAD_DIR.mkdir(exist_ok=True)

# Q&A ENDPOINT
@app.post("/ask", response_model=AnswerResponse)
def ask_question(req: QuestionRequest):
    if not store:
        return AnswerResponse(
            question=req.question,
            answer="Vector database not loaded",
            sources_used=0
        )
    qvec = embed_query(req.question)
    results = store.search(qvec, k=req.k)
    answer = answer_question_with_llm(req.question, results)
    
    return AnswerResponse(
        question=req.question,
        answer=answer,
        sources_used=len(results)
    )

# REPORT GENERATION ENDPOINT
@app.post("/generate_report", response_model=ReportResponse)
def generate_report(req: ReportRequest):
    if not store:
        return ReportResponse(
            message="Vector database not loaded",
            filename="none",
            download_url="empty"
        )
    report_data = build_report_from_query(
        store=store,
        query=req.query,
        k=req.k
    )
    output_dir = Path("outputs/reports")
    output_dir.mkdir(parents=True, exist_ok=True)
    report_path = output_dir / req.filename
    render_docx(report_data, str(report_path))
    
    return ReportResponse(
        message="✅ Report generated successfully",
        filename=req.filename,
        download_url=f"/download/{req.filename}"
    )

# FILE UPLOAD & INGESTION
@app.post("/upload")
def upload_file(file: UploadFile = File(...)):
    file_path = UPLOAD_DIR / file.filename
    
    # Save uploaded file
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    # Run incremental ingestion
    ingest_folder_incremental(str(UPLOAD_DIR))
    
    # Reload FAISS store
    global store
    store = FaissStore.load(config.INDEX_DIR)
    
    return {
        "message": "✅ File uploaded and indexed successfully",
        "filename": file.filename
    }

# LIST ALL UPLOADED FILES
@app.get("/files")
def list_uploaded_files():
    files = [f.name for f in UPLOAD_DIR.iterdir() if f.is_file()]
    return {"files": files}

# DOWNLOAD ENDPOINT
@app.get("/download/{filename}")
def download_report(filename: str):
    file_path = Path("outputs/reports") / filename
    if not file_path.exists():
        return {"error": "File not found"}
    
    return FileResponse(
        path=file_path,
        filename=filename,
        media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    )

# HEALTH CHECK
@app.get("/status")
def status():
    if store:
        return {
            "status": "Ready",
            "chunks": len(store.meta),
            "message": "Vector database loaded. You can upload files, ask questions or generate reports."
        }
    else:
        return {
            "status": "Not Ready",
            "message": "Vector database not loaded. Upload PDFs first."
        }

@app.get("/")
def home():
    return {"welcome": "Hello"}