main.py

from fastapi import FastAPI, UploadFile, File, Form, Depends, HTTPException, Request, Body
from fastapi.responses import JSONResponse
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from pydantic import BaseModel, EmailStr, validator
import uuid
import re

from config import settings
from services.ingestion import ingest_document
from services.rag import rag_query
from database.models import Base, Booking
from vectorstore.pinecone_client import index, get_namespace
from fastapi.middleware.cors import CORSMiddleware


frontend = settings.FRONTEND_URL



# ====================== DATABASE ======================
engine = create_async_engine(settings.DATABASE_URL, echo=False)
async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

async def get_db():
    async with async_session() as session:
        yield session

# ====================== APP ======================
app = FastAPI(
    title="Custom RAG",
    description="Upload PDF → Chat → Book → New Chat",
    version="2.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[frontend, "http://localhost:5173"],  # Use localhost only
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ====================== BOOKING MODEL ======================
class BookingRequest(BaseModel):
    name: str
    email: EmailStr
    date: str
    time: str

    @validator("date")
    def validate_date(cls, v):
        if not re.match(r"^\d{4}-\d{2}-\d{2}$", v):
            raise ValueError("Date must be YYYY-MM-DD")
        return v

    @validator("time")
    def validate_time(cls, v):
        if not re.match(r"^([0-1]?[0-9]|2[0-3]):[0-5][0-9]$", v):
            raise ValueError("Time must be HH:MM (24h)")
        return v

# ====================== STARTUP ======================
@app.on_event("startup")
async def startup():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

# ====================== ENDPOINTS ======================

@app.post("/upload")
async def upload_document(
    request: Request,
    file: UploadFile = File(...),
    chunking_strategy: str = Form("recursive"),
    db: AsyncSession = Depends(get_db)
):
    print("="*50)
    print("[UPLOAD] Request received")
    print(f"[UPLOAD] All cookies: {dict(request.cookies)}")
    
    if not file.filename.lower().endswith((".pdf", ".txt")):
        raise HTTPException(400, detail="Only PDF/TXT allowed")

    # Get or create session
    session_id = request.cookies.get("session_id")
    if not session_id:
        session_id = str(uuid.uuid4())
        print(f"[UPLOAD] Created NEW session: {session_id}")
    else:
        print(f"[UPLOAD] Using EXISTING session: {session_id}")
    
    result = await ingest_document(file=file, strategy=chunking_strategy, session_id=session_id, db=db)

    # Create response and set cookie
    response = JSONResponse({
        "status": "success",
        "chunks": result["chunks"],
        "message": "PDF uploaded successfully!",
        "session_id": session_id  # Send back for debugging
    })
    
    response.set_cookie(
        key="session_id",
        value=session_id,
        httponly=False,  # Allow JS to read
        max_age=86400,
        samesite="lax",
        secure=False,
        path="/",
        domain=None  # Don't set domain, let browser handle it
    )
    
    print(f"[UPLOAD] Response: Setting cookie session_id={session_id}")
    print("="*50)
    return response


@app.post("/chat")
async def chat(
    request: Request,
    message: str = Body(..., embed=True),
    db: AsyncSession = Depends(get_db)
):
    print("="*50)
    print("[CHAT] Request received")
    print(f"[CHAT] All cookies: {dict(request.cookies)}")
    
    if not message.strip():
        raise HTTPException(400, detail="Message required")

    # Get or create session
    session_id = request.cookies.get("session_id")
    if not session_id:
        session_id = str(uuid.uuid4())
        print(f"[CHAT] Created NEW session: {session_id}")
    else:
        print(f"[CHAT] Using EXISTING session: {session_id}")
    
    bot_response = await rag_query(session_id, message.strip())
    
    # Create response and set cookie
    response = JSONResponse({
        "response": bot_response,
        "session_id": session_id  # Send back for debugging
    })
    
    response.set_cookie(
        key="session_id",
        value=session_id,
        httponly=False,
        max_age=86400,
        samesite="lax",
        secure=False,
        path="/",
        domain=None
    )
    
    print(f"[CHAT] Response: Setting cookie session_id={session_id}")
    print("="*50)
    return response


@app.post("/book-interview")
async def book_interview(
    request: Request,
    booking: BookingRequest,
    db: AsyncSession = Depends(get_db)
):
    print("="*50)
    print("[BOOKING] Request received")
    print(f"[BOOKING] All cookies: {dict(request.cookies)}")
    
    session_id = request.cookies.get("session_id")
    if not session_id:
        session_id = str(uuid.uuid4())
        print(f"[BOOKING] Created NEW session: {session_id}")
    else:
        print(f"[BOOKING] Using EXISTING session: {session_id}")
    
    new_booking = Booking(
        name=booking.name,
        email=booking.email,
        date=booking.date,
        time=booking.time
    )
    db.add(new_booking)
    await db.commit()

    response = JSONResponse({
        "status": "success",
        "message": f"Interview booked for {booking.name} on {booking.date} at {booking.time}!",
        "session_id": session_id
    })
    
    response.set_cookie(
        key="session_id",
        value=session_id,
        httponly=False,
        max_age=86400,
        samesite="lax",
        secure=False,
        path="/"
    )
    
    print(f"[BOOKING] Response: Setting cookie session_id={session_id}")
    print("="*50)
    return response


@app.post("/new-chat")
async def new_chat(request: Request):
    print("="*50)
    print("[NEW-CHAT] Request received")
    print(f"[NEW-CHAT] All cookies: {dict(request.cookies)}")
    
    session_id = request.cookies.get("session_id")
    
    if not session_id:
        print("[NEW-CHAT] No session found")
        return JSONResponse({"message": "No active session"})

    print(f"[NEW-CHAT] Clearing session: {session_id}")
    namespace = get_namespace(session_id)

    try:
        index.delete(namespace=namespace, delete_all=True)
        print(f"[NEW-CHAT] Cleared Pinecone namespace: {namespace}")
    except Exception as e:
        print(f"[NEW-CHAT] Pinecone cleanup failed: {e}")

    # Generate new session ID
    new_session_id = str(uuid.uuid4())
    print(f"[NEW-CHAT] New session ID: {new_session_id}")
    
    response = JSONResponse({
        "message": "New chat started! Previous document removed.",
        "status": "ready_for_new_document",
        "session_id": new_session_id
    })
    
    response.set_cookie(
        key="session_id",
        value=new_session_id,
        httponly=False,
        max_age=86400,
        samesite="lax",
        secure=False,
        path="/",
        domain=None
    )
    
    print(f"[NEW-CHAT] Response: Setting cookie session_id={new_session_id}")
    print("="*50)
    return response


@app.get("/")
async def root():
    return {"message": "RAG API is running! Visit /docs"}