# server.py
import json
import traceback
import uuid
from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
import pycountry
from pydantic import BaseModel
from chat_utils 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):
    language: str
    text: str
    session_id: str | None = None  # Optional session ID from client

LANG_NATIVE_NAMES = {
    "en": "English",
    "fr": "Français",
    "es": "Español",
    "hi": "हिन्दी",
    "bn": "বাংলা",
    "te": "తెలుగు",
    "mr": "मराठी",
    "ta": "தமிழ்",
    "ur": "اردو",
    "gu": "ગુજરાતી",
    "kn": "ಕನ್ನಡ",
    "ml": "മലയാളം",
    "pa": "ਪੰਜਾਬੀ",
    "as": "অসমীয়া",
    "mai": "मैथिली",
    "sd": "سنڌي",
    "sat": "ᱥᱟᱱᱛᱟᱲᱤ",
}

@router.get("/languages")
async def handle_fetch_languages():
    supported_lang_codes = [
        "en","fr","es","hi","bn","te","mr","ta","ur",
        "gu","kn","ml","pa","as","mai","sd","sat",
    ]

    languages = []
    for code in supported_lang_codes:
        lang = pycountry.languages.get(alpha_2=code) or pycountry.languages.get(alpha_3=code)
        if lang is None:
            continue  # skip unknown codes
        english_name = lang.name
        native_name = LANG_NATIVE_NAMES.get(code, english_name)
        languages.append({
            "code": code,
            "name": english_name,
            "native_name": native_name,
        })

    languages.sort(key=lambda x: x["name"])
    return languages


@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_units = SanatanDatabase().count(
            collection_name=scripture["collection_name"]
        )
        markdown += f"- {scripture['title']} : `{num_units}` {scripture["unit"]}s\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())
        print(session_id, ": user sent message : ", msg.text)
        # 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,
            preferred_language=msg.language or "English",
        )

        # 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}"})