# server.py
import json
import random
import traceback
from typing import Optional
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
from modules.quiz.answer_validator import validate_answer
from modules.quiz.models import Question
from modules.quiz.quiz_helper import generate_question

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


class QuizGeneratePayload(BaseModel):
    language: Optional[str] = "English"
    scripture: Optional[str] = None
    complexity: Optional[str] = None
    mode: Optional[str] = None
    session_id: Optional[str] = None  # Optional session ID from client


class QuizEvalPayload(BaseModel):
    language: Optional[str] = "English"
    q: Question
    answer: str
    session_id: Optional[str] = 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}"})


@router.post("/quiz/generate")
async def handle_quiz_generate(payload: QuizGeneratePayload, request: Request):
    q = generate_question(
        collection=payload.scripture
        or random.choice(
            [
                s["collection_name"]
                for s in SanatanConfig.scriptures
                if s["collection_name"] != "yt_metadata"
            ]
        ),
        complexity=payload.complexity
        or random.choice(["beginner", "intermediate", "advanced"]),
        mode=payload.mode or random.choice(["mcq", "open"]),
        preferred_lamguage=payload.language or "English",
    )
    print(q.model_dump_json(indent=1))
    return q.model_dump()


@router.post("/quiz/eval")
async def handle_quiz_eval(payload: QuizEvalPayload, request: Request):
    result = validate_answer(
        payload.q, payload.answer, preferred_language=payload.language or "English"
    )
    return result

@router.get("/scriptures")
async def handle_get_scriptures():
    return_values = {}
    for scripture in SanatanConfig().scriptures:
        num_units = SanatanDatabase().count(
            collection_name=scripture["collection_name"]
        )
        return_values[scripture['title']] = scripture['collection_name']
    return return_values