import os
import gradio as gr
from openai import OpenAI

# ----------------------------------------------------------------------
# Helper to read a secret (fallback is useful when you run locally)
# ----------------------------------------------------------------------
def _secret(key: str, fallback: str = "") -> str:
    """Return the value of a secret or the supplied fallback."""
    return os.getenv(key, fallback)


# ----------------------------------------------------------------------
# Core chat logic – system prompt comes from the secret `prec_chat`
# ----------------------------------------------------------------------
def respond(
    message: str,
    history: list[dict[str, str]],
    # model_name: str,  
):
    """
    Generate a response using OpenRouter API via OpenAI client.
    
    * System prompt = secret `prec_chat`
    * OpenRouter API key = secret `OPENROUTER_API_KEY`
    """
    # 1️⃣ Load the system prompt (fallback = generic assistant)
    # system_message = _secret("prec_chat", "You are a helpful assistant.")

    # 2️⃣ Load the OpenRouter API key
    openrouter_api_key = _secret("OPENROUTER_API_KEY")
    if not openrouter_api_key:
        raise RuntimeError(
            "OPENROUTER_API_KEY not found in secrets. Add it to secrets.toml (or via the Space UI)."
        )

    # 3️⃣ Initialize OpenAI client with OpenRouter configuration
    # client = OpenAI(
    #     base_url="https://openrouter.ai/api/v1",        
    #     api_key=openrouter_api_key,
    # )

    client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=openrouter_api_key,
    default_headers={
        # "HTTP-Referer": "<YOUR_SITE_URL>",  # Optional. Set your site URL for tracking.
        "X-Title": "hf_precise",        # Optional. Set your app name for tracking.
        "X-OpenRouter-Preset": "@preset/precise-chat-agent"  # <<< The crucial line
    },
)

    # 4️⃣ Build the message list for the chat completion
    messages = []
    messages.extend(history)                     # previous conversation turns
    messages.append({"role": "user", "content": message})  # current user query

    # 5️⃣ Stream the response back to the UI
    response = ""
    stream = client.chat.completions.create(
        # model=model_name,
        messages=messages,
        max_tokens=8096,
        stream=True
    )
    
    for chunk in stream:
        if chunk.choices[0].delta.content is not None:
            token = chunk.choices[0].delta.content
            response += token
            yield response


# # ──────────────────────────────────────────────────────────────────────
# # List of models available through OpenRouter
# # ──────────────────────────────────────────────────────────────────────
# AVAILABLE_MODELS = [
# "@preset/precise-chat-agent"
# ]

# ----------------------------------------------------------------------
# UI – the system‑prompt textbox has been removed.
# ----------------------------------------------------------------------
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    # additional_inputs=[
    #     gr.Dropdown(
    #         choices=AVAILABLE_MODELS,
    #         value=AVAILABLE_MODELS[0],
    #         label="Model",
    #         interactive=True,
    #     ),        
    # ],
)

# ----------------------------------------------------------------------
# Assemble the Blocks layout (no LoginButton – we use basic auth)
# ----------------------------------------------------------------------
with gr.Blocks() as demo:
    chatbot.render()


# ----------------------------------------------------------------------
# Launch – protect the UI with the credentials from secrets.
# ----------------------------------------------------------------------
if __name__ == "__main__":
    # ------------------------------------------------------------------
    # 1️⃣ Pull the allowed credentials from secrets (fail fast if missing)
    # ------------------------------------------------------------------
    allowed_user = _secret("CHAT_USER")
    allowed_pass = _secret("CHAT_PASS")

    if not allowed_user or not allowed_pass:
        raise RuntimeError(
            "Authentication credentials not found in secrets. "
            "Add CHAT_USER and CHAT_PASS to secrets.toml (or via the HF Spaces UI)."
        )

    # ------------------------------------------------------------------
    # 2️⃣ Launch
    # ------------------------------------------------------------------
    demo.launch(
        auth=(allowed_user, allowed_pass),   # <-- Gradio's built‑in basic auth
        ssr_mode=False,                     # <-- avoids the i18n locale error
        # In a Space we **must not** set share=True (Spaces already give a public URL)
        # If you run locally and want a shareable link, add share=True here.
        server_name="0.0.0.0",              # listen on all interfaces (needed in containers)
        # Optional: give the app a nice title
        # title="Secure Chatbot",
    )