# server.py
import traceback
import uuid
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from app import chat
from config import SanatanConfig
from db import SanatanDatabase

router = APIRouter()

# In-memory mapping from session_id -> thread_id
# For production, you may want Redis or a DB for persistence
thread_map = {}

class Message(BaseModel):
    text: str
    session_id: str | None = None  # Optional session ID from client

@router.post("/greet")
async def handle_greet(msg: Message):
    markdown = "Namaskaram 🙏 I am **bhashyam.ai** and I can help you explore the following scriptures:\n---\n"
    for scripture in SanatanConfig().scriptures:
        num_records = SanatanDatabase().count(collection_name=scripture["collection_name"])
        markdown += f"- {scripture['title']} : `{num_records}` verses\n"

    session_id = msg.session_id
    if not session_id:
        session_id = str(uuid.uuid4())
    return {"reply": markdown, "session_id": session_id}

@router.post("/chat")
async def handle_chat(msg: Message, request: Request):
    try:
        # Use existing session_id if provided, else generate new
        session_id = msg.session_id
        if not session_id:
            session_id = str(uuid.uuid4())

        # Get or create a persistent thread_id for this session
        if session_id not in thread_map:
            thread_map[session_id] = str(uuid.uuid4())
        thread_id = thread_map[session_id]

        # Call your graph/chat function
        reply_text = chat(
            debug_mode=False,
            message=msg.text,
            history=None,
            thread_id=thread_id
        )

        # Return both reply and session_id to the client
        return {"reply": reply_text, "session_id": session_id}

    except Exception as e:
        traceback.print_exc()
        return JSONResponse(
            status_code=500,
            content={"reply": f"Error: {e}"}
        )