app/main.py

"""FastAPI application — main entry point with streaming support."""

import os
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional

from app.session import session_manager
from app.graph import run_chat, run_chat_streaming


# --- App Setup ---
app = FastAPI(
    title="StackLogix Chatbot",
    description="Intelligent chatbot powered by LangGraph + Qdrant",
    version="1.0.0",
)

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

# --- Static Files ---
static_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "static")
if os.path.exists(static_dir):
    app.mount("/static", StaticFiles(directory=static_dir), name="static")


# --- Request / Response Models ---
class ChatRequest(BaseModel):
    session_id: str
    user_message: str


class ChatResponse(BaseModel):
    answer: str
    follow_up_question: Optional[str] = None


class IngestResponse(BaseModel):
    status: str
    message: str


# --- Endpoints ---

@app.get("/")
async def serve_ui():
    """Serve the chat UI."""
    index_path = os.path.join(static_dir, "index.html")
    if os.path.exists(index_path):
        return FileResponse(index_path)
    return {"message": "StackLogix Chatbot API is running. No UI found in /static."}


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "ok", "service": "stacklogix-chatbot"}


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """Non-streaming chat endpoint (backward compatible)."""
    if not request.user_message.strip():
        raise HTTPException(status_code=400, detail="user_message cannot be empty")

    if not request.session_id.strip():
        raise HTTPException(status_code=400, detail="session_id cannot be empty")

    try:
        session = session_manager.get_or_create(request.session_id)
        result = run_chat(session, request.user_message.strip())

        return ChatResponse(
            answer=result["answer"],
            follow_up_question=result.get("follow_up_question"),
        )

    except Exception as e:
        print(f"Chat error: {e}")
        raise HTTPException(
            status_code=500,
            detail="An error occurred while processing your message. Please try again.",
        )


@app.post("/chat/stream")
async def chat_stream(request: ChatRequest):
    """
    Streaming chat endpoint — returns Server-Sent Events (SSE).
    Events:
      data: {"type": "token", "content": "..."}
      data: {"type": "sources", "sources": [...]}
      data: {"type": "follow_up", "content": "..."}
      data: {"type": "done"}
    """
    if not request.user_message.strip():
        raise HTTPException(status_code=400, detail="user_message cannot be empty")

    if not request.session_id.strip():
        raise HTTPException(status_code=400, detail="session_id cannot be empty")

    session = session_manager.get_or_create(request.session_id)

    def event_generator():
        try:
            yield from run_chat_streaming(session, request.user_message.strip())
        except Exception as e:
            import json
            print(f"Stream error: {e}")
            yield f"data: {json.dumps({'type': 'error', 'content': 'An error occurred. Please try again.'})}\n\n"
            yield f"data: {json.dumps({'type': 'done'})}\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("/ingest", response_model=IngestResponse)
async def ingest():
    """Trigger document ingestion pipeline."""
    try:
        from app.ingestion import ingest_all_documents
        ingest_all_documents()
        return IngestResponse(status="success", message="Documents ingested successfully.")
    except Exception as e:
        print(f"Ingestion error: {e}")
        raise HTTPException(status_code=500, detail=f"Ingestion failed: {str(e)}")


@app.get("/collection-info")
async def collection_info():
    """Get Qdrant collection information."""
    try:
        from app.retriever import get_retriever
        retriever = get_retriever()
        info = retriever.get_collection_info()
        return info
    except Exception as e:
        return {"error": str(e)}


@app.get("/chats")
async def list_chats():
    """List all chat sessions with summary info."""
    from app.chat_db import get_all_sessions
    sessions = get_all_sessions()
    return {"sessions": sessions, "total": len(sessions)}


@app.get("/chats/{session_id}")
async def get_chat(session_id: str):
    """Get full chat history for a specific session."""
    from app.chat_db import get_chat_history
    history = get_chat_history(session_id)
    if not history:
        raise HTTPException(status_code=404, detail=f"No chat history found for session: {session_id}")
    return {
        "session_id": session_id,
        "messages": history,
        "message_count": len(history),
    }