# RAG Admin Panel - PDF Upload Management (Port 9000)
import os
from pathlib import Path
from contextlib import asynccontextmanager

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.requests import Request
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware

# Import RAG module
from app import rag

IS_HF_SPACE = bool(os.getenv("SPACE_ID"))

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Ensure RAG index is ready when admin panel starts."""
    if not IS_HF_SPACE:
        loaded = rag.load_vector_store()
        if not loaded:
            rag.rebuild_vector_store_from_pdfs()
    yield


# Initialize FastAPI app for admin
admin_app = FastAPI(title="RAG Admin Panel", version="1.0.0", lifespan=lifespan)

# Add CORS middleware
admin_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files and templates
BASE_DIR = Path(__file__).resolve().parent
admin_app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static")
templates = Jinja2Templates(directory=BASE_DIR / "templates")


@admin_app.get("/", response_class=HTMLResponse)
async def admin_home(request: Request):
    """Render the admin panel for PDF upload"""
    return templates.TemplateResponse("admin.html", {"request": request})


@admin_app.post("/api/upload")
async def upload_pdf(file: UploadFile = File(...)):
    """Upload a PDF file for RAG processing"""
    if not file.filename.lower().endswith('.pdf'):
        raise HTTPException(status_code=400, detail="Only PDF files are allowed")
    
    try:
        # Initialize embeddings if not already done
        rag.initialize_embeddings()
        
        # Save uploaded file
        RAG_DATA_DIR = Path(__file__).resolve().parent.parent / "rag_data"
        RAG_DATA_DIR.mkdir(parents=True, exist_ok=True)
        
        pdf_path = RAG_DATA_DIR / file.filename
        
        content = await file.read()
        with open(pdf_path, "wb") as f:
            f.write(content)
        
        # Process the PDF
        chunks = rag.load_and_process_pdf(str(pdf_path))
        
        if not chunks:
            raise HTTPException(status_code=400, detail="Could not extract text from PDF")
        
        # Create/update vector store
        success = rag.create_vector_store(chunks)
        
        if success:
            rag.get_rag_status()
            return JSONResponse({
                "success": True,
                "message": f"PDF '{file.filename}' uploaded and processed successfully",
                "chunks_created": len(chunks),
                "total_documents": len(rag.uploaded_documents)
            })
        else:
            raise HTTPException(status_code=500, detail="Failed to create vector store")
            
    except Exception as e:
        print(f"RAG Upload Error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Failed to process PDF: {str(e)}")


@admin_app.get("/api/status")
async def get_status():
    """Get RAG system status"""
    return JSONResponse(rag.get_rag_status())


@admin_app.post("/api/clear")
async def clear_data():
    """Clear all RAG data"""
    rag.clear_rag_data()
    return JSONResponse({"success": True, "message": "RAG data cleared"})


@admin_app.delete("/api/document/{filename}")
async def delete_document(filename: str):
    """Delete a specific document"""
    try:
        RAG_DATA_DIR = Path(__file__).resolve().parent.parent / "rag_data"
        pdf_path = RAG_DATA_DIR / filename
        
        if pdf_path.exists():
            os.remove(pdf_path)
        
        if list(RAG_DATA_DIR.glob("*.pdf")):
            rag.rebuild_vector_store_from_pdfs()
        else:
            rag.clear_rag_data()
        
        return JSONResponse({"success": True, "message": f"Document '{filename}' deleted"})
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to delete: {str(e)}")


@admin_app.post("/api/rebuild")
async def rebuild_data():
    """Rebuild vector store from all PDFs in rag_data."""
    success = rag.rebuild_vector_store_from_pdfs()
    if success:
        return JSONResponse({"success": True, "message": "RAG rebuilt successfully from all PDFs"})
    return JSONResponse({"success": False, "message": "No valid PDFs found to rebuild RAG"})


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(admin_app, host="0.0.0.0", port=9000)