app/main.py

import json
import logging
import uuid
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

from app.config import settings
from app.models.schemas import (
    FinalResponse,
    QueryRequest,
    ScrapeRequest,
    StatusResponse,
    SummarizeRequest,
    SummaryResponse,
)
from app.services.rag_service import RAGService
from app.utils.document_processor import ALLOWED_EXTENSIONS, DocumentProcessor

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("studyson")


@asynccontextmanager
async def lifespan(_: FastAPI):
    settings.upload_dir.mkdir(parents=True, exist_ok=True)
    settings.chroma_dir.mkdir(parents=True, exist_ok=True)
    logger.info("Studyson starting | model=%s | chroma=%s", settings.groq_model, settings.chroma_dir)
    yield


app = FastAPI(
    title="Studyson RAG API",
    description="Document QA and summarization using RAG (Groq + LlamaIndex + Chroma)",
    version="2.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

rag_service = RAGService()
doc_processor = DocumentProcessor()

app.mount("/static", StaticFiles(directory="static"), name="static")


@app.get("/")
async def read_root():
    return FileResponse("static/index.html")


@app.post("/upload", response_model=StatusResponse)
async def upload_document(file: UploadFile = File(...)):
    if not file.filename or not doc_processor.validate_file_type(file.filename):
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type. Allowed: {', '.join(sorted(ALLOWED_EXTENSIONS))}",
        )

    if file.size and file.size > settings.max_file_size:
        raise HTTPException(
            status_code=400,
            detail=f"File exceeds {settings.max_file_size // (1024 * 1024)} MB limit",
        )

    settings.upload_dir.mkdir(parents=True, exist_ok=True)
    safe_name = file.filename.replace("/", "_").replace("\\", "_")
    file_path = settings.upload_dir / safe_name

    try:
        content = await file.read()
        if len(content) > settings.max_file_size:
            raise HTTPException(status_code=400, detail="File exceeds size limit")
        file_path.write_bytes(content)

        text = await doc_processor.extract_text(file_path)
        cleaned_text = doc_processor.clean_text(text)
        if not cleaned_text.strip():
            raise HTTPException(status_code=400, detail="No extractable text in file")

        rag_service.add_document(cleaned_text, safe_name)

        return StatusResponse(
            status="success",
            message=f"Document '{safe_name}' indexed successfully",
            details={
                "filename": safe_name,
                "text_length": len(cleaned_text),
                "indexed_documents": rag_service.get_indexed_documents(),
            },
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Upload failed for %s", file.filename)
        if file_path.exists():
            file_path.unlink(missing_ok=True)
        raise HTTPException(status_code=500, detail=f"Error processing document: {e}")


@app.post("/scrape_and_index", response_model=StatusResponse)
async def scrape_and_index(request: ScrapeRequest):
    try:
        title, text = await doc_processor.scrape_url(str(request.url))
        cleaned_text = doc_processor.clean_text(text)
        if not cleaned_text.strip():
            raise HTTPException(status_code=400, detail="No extractable text on page")

        rag_service.add_document(cleaned_text, title)

        return StatusResponse(
            status="success",
            message="URL content indexed successfully",
            details={
                "url": str(request.url),
                "title": title,
                "text_length": len(cleaned_text),
                "indexed_documents": rag_service.get_indexed_documents(),
            },
        )
    except HTTPException:
        raise
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.exception("Scrape failed for %s", request.url)
        raise HTTPException(status_code=500, detail=f"Error scraping URL: {e}")


def _resolve_session(session_id: str | None) -> str:
    return session_id or str(uuid.uuid4())


@app.post("/stream_query")
async def stream_query(request: QueryRequest):
    if not rag_service.has_documents():
        raise HTTPException(status_code=400, detail="No documents indexed. Please upload a document first.")

    session_id = _resolve_session(request.session_id)

    async def event_generator() -> AsyncGenerator[str, None]:
        try:
            yield f"data: {json.dumps({'session_id': session_id})}\n\n"

            answer_parts: list[str] = []
            async for token in rag_service.stream_query(request.question, session_id):
                answer_parts.append(token)
                yield f"data: {json.dumps({'token': token})}\n\n"

            full_answer = "".join(answer_parts)
            _, sources = await rag_service.query(request.question)

            final = FinalResponse(
                final_answer=full_answer,
                sources=[s.model_dump() for s in sources],
            )
            yield "data: [DONE]\n\n"
            yield f"data: {json.dumps(final.model_dump())}\n\n"
        except Exception as e:
            logger.exception("Stream query failed")
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.post("/query", response_model=FinalResponse)
async def query(request: QueryRequest):
    if not rag_service.has_documents():
        raise HTTPException(status_code=400, detail="No documents indexed. Please upload a document first.")
    try:
        answer, sources = await rag_service.query(request.question)
        return FinalResponse(final_answer=answer, sources=sources)
    except Exception as e:
        logger.exception("Query failed")
        raise HTTPException(status_code=500, detail=f"Error processing query: {e}")


@app.post("/summarize", response_model=SummaryResponse)
async def summarize(request: SummarizeRequest):
    if not rag_service.has_documents():
        raise HTTPException(status_code=400, detail="No documents indexed. Please upload a document first.")
    try:
        summary = await rag_service.summarize(max_length=request.max_length)
        return SummaryResponse(
            summary=summary,
            word_count=len(summary.split()),
            source_documents=rag_service.get_indexed_documents(),
        )
    except Exception as e:
        logger.exception("Summarize failed")
        raise HTTPException(status_code=500, detail=f"Error generating summary: {e}")


@app.post("/reset", response_model=StatusResponse)
async def reset_index():
    try:
        rag_service.reset_all()
        if settings.upload_dir.exists():
            for path in settings.upload_dir.glob("*"):
                if path.is_file():
                    path.unlink(missing_ok=True)
        return StatusResponse(status="success", message="Index reset. All documents removed.")
    except Exception as e:
        logger.exception("Reset failed")
        raise HTTPException(status_code=500, detail=f"Error resetting index: {e}")


@app.get("/status", response_model=StatusResponse)
async def get_status():
    docs = rag_service.get_indexed_documents()
    return StatusResponse(
        status="online",
        message="Studyson RAG API is running",
        details={
            "model": settings.groq_model,
            "has_documents": rag_service.has_documents(),
            "indexed_documents": docs,
            "document_count": len(docs),
        },
    )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app.main:app", host=settings.host, port=settings.port)