import os
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from pydantic import BaseModel, Field
from dotenv import load_dotenv
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from src.new_graph import build_agent, AgentContext
from src.embedding import upload_file
import asyncio
from psycopg_pool import AsyncConnectionPool
import logging
import traceback
logger = logging.getLogger("pdf_rag_api")
logging.basicConfig(level=logging.INFO)

load_dotenv()

DB_URI = os.getenv("DATABASE_URI")

@asynccontextmanager
async def lifespan(app: FastAPI):
    async with AsyncConnectionPool(
        conninfo=DB_URI,
        min_size=1,
        max_size=15,
        max_idle=300,          
        max_lifetime=1800,    
        reconnect_timeout=10,
        kwargs={"autocommit": True, "prepare_threshold": 0},
        check=AsyncConnectionPool.check_connection,
        open=False,            
    ) as pool:
        await pool.open(wait=True, timeout=15)
        checkpointer = AsyncPostgresSaver(pool)
        await checkpointer.setup()
        app.state.pool = pool
        app.state.agent = build_agent(checkpointer)
        yield
    
app = FastAPI(
    title="Enterprise PDF RAG API",
    description="A production-grade backend powering an agentic LangChain RAG assistant.",
    version="2.0.0",
    lifespan=lifespan
)

class ChatRequest(BaseModel):
    message: str = Field(...)
    user_id: str = Field(...)
    thread_id: str = Field(...)

class ChatResponse(BaseModel):
    status: str
    thread_id: str
    response: str

@app.post("/chat",response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
    agent = app.state.agent  

    config = {"configurable": {"thread_id": request.thread_id}}

    try:
        result = await agent.ainvoke(
            {"messages": [{"role": "user", "content": request.message}]},
            config=config,
            context=AgentContext(user_id=request.user_id),
        )
    except Exception as e:
        logger.error("Agent processing failed: %s", e)
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail="Agent Processing Error!")
    output_messages = result.get("messages", [])
    if not output_messages:
        raise HTTPException(status_code=500, detail="No messages returned from the agent.")

    return ChatResponse(
        status="success",
        thread_id=request.thread_id,
        response=output_messages[-1].content,
    )

@app.post("/upload")
async def upload_pdf(file: UploadFile = File(...),user_id: str = Form(..., min_length=1) ):
    if file.content_type != "application/pdf":
        raise HTTPException(status_code=400, detail="Only PDF files are supported.")
    
    file_bytes = await file.read()
    if not file_bytes:
        raise HTTPException(status_code=400, detail="Uploaded file is empty.")


    await asyncio.to_thread(upload_file,file_bytes,file.filename,user_id)

    return {
        "status": "success",
        "message": "Document processed. You can now ask questions."
    }