import os

import gradio as gr
import requests
from models import MODEL_IDS

# Modal endpoint URL - set this after deploying backend.py
MODAL_ENDPOINT = os.environ.get("MODAL_ENDPOINT", "")
# API key for authenticating with Modal backend
MODEL_SITE_API_KEY = os.environ.get("MODEL_SITE_API_KEY", "")
# Password for Gradio login (any username accepted)
SITE_PASSWORD = os.environ.get("SITE_PASSWORD", "")


def make_respond_fn(model_id: str):
    def respond(message: str, history: list[tuple[str, str]]) -> str:
        if not MODAL_ENDPOINT:
            return "Error: MODAL_ENDPOINT environment variable not set"

        try:
            response = requests.post(
                MODAL_ENDPOINT,
                headers={"X-API-Key": MODEL_SITE_API_KEY},
                json={
                    "model_id": model_id,
                    "message": message,
                    "history": history,
                },
                timeout=120,  # Cold start can take a while
            )
            response.raise_for_status()
            data = response.json()

            if "error" in data:
                return f"Error: {data['error']}"

            return data.get("response", "No response received")

        except requests.exceptions.Timeout:
            return "Error: Request timed out. The model may be starting up, please try again."
        except requests.exceptions.RequestException as e:
            return f"Error: {e}"

    return respond


# Create tabbed interface with one chat per model
with gr.Blocks(title="posttraining-practice") as demo:
    gr.Markdown("# posttraining-practice")
    gr.Markdown("Chat with different fine-tuned models")

    missing = [
        v
        for v in ["MODAL_ENDPOINT", "MODEL_SITE_API_KEY", "SITE_PASSWORD"]
        if not os.environ.get(v)
    ]
    if missing:
        gr.Markdown(f"⚠️ **Warning:** Missing secrets: {', '.join(missing)}")

    with gr.Tabs():
        for model_id in MODEL_IDS:
            short_name = model_id.split("/")[-1]
            with gr.Tab(short_name):
                gr.ChatInterface(
                    fn=make_respond_fn(model_id),
                    description=f"Chatting with: {model_id}",
                )


def check_password(username: str, password: str) -> bool:
    return password == SITE_PASSWORD


if __name__ == "__main__":
    auth = check_password if SITE_PASSWORD else None
    demo.launch(auth=auth)