from __future__ import annotations

import json
import time
from typing import Any

import gradio as gr
import requests


BASE_URL = "https://api.velokey.ai/v1"
CHAT_COMPLETIONS_URL = f"{BASE_URL}/chat/completions"
MODELS_URL = f"{BASE_URL}/models"

MODEL_EXAMPLES = [
    "gpt-5.5",
    "claude-sonnet-4-6",
    "gemini-3-pro-preview",
    "deepseek-v4-pro",
    "qwen3.7-max",
]

DEFAULT_SYSTEM_PROMPT = "You are a concise assistant for developers."
DEFAULT_USER_PROMPT = "Explain what an OpenAI-compatible API gateway is in two sentences."


def _headers(api_key: str) -> dict[str, str]:
    return {
        "Authorization": f"Bearer {api_key.strip()}",
        "Content-Type": "application/json",
        "User-Agent": "velokey-huggingface-playground/1.0",
    }


def _format_json(data: Any) -> str:
    return json.dumps(data, ensure_ascii=False, indent=2)


def _request_json(method: str, url: str, api_key: str, **kwargs: Any) -> tuple[int, Any, float]:
    start = time.perf_counter()
    response = requests.request(
        method,
        url,
        headers=_headers(api_key),
        timeout=60,
        **kwargs,
    )
    elapsed_ms = (time.perf_counter() - start) * 1000

    try:
        body: Any = response.json()
    except ValueError:
        body = response.text

    return response.status_code, body, elapsed_ms


def list_models(api_key: str) -> tuple[str, str]:
    if not api_key.strip():
        return "Paste a VeloKey API key first.", ""

    try:
        status, body, elapsed_ms = _request_json("GET", MODELS_URL, api_key)
    except requests.RequestException as exc:
        return f"Request failed: {exc}", ""

    if status >= 400:
        return f"Model list request returned HTTP {status}.", _format_json(body)

    model_ids: list[str] = []
    if isinstance(body, dict) and isinstance(body.get("data"), list):
        for item in body["data"]:
            if isinstance(item, dict) and item.get("id"):
                model_ids.append(str(item["id"]))
    elif isinstance(body, list):
        for item in body:
            if isinstance(item, dict) and item.get("id"):
                model_ids.append(str(item["id"]))

    if model_ids:
        preview = "\n".join(model_ids[:30])
        if len(model_ids) > 30:
            preview += f"\n...and {len(model_ids) - 30} more"
        summary = f"Found {len(model_ids)} model IDs in {elapsed_ms:.0f} ms."
        return summary, preview

    return f"Request succeeded in {elapsed_ms:.0f} ms, but no model IDs were recognized.", _format_json(body)


def run_chat_completion(
    api_key: str,
    model: str,
    system_prompt: str,
    user_prompt: str,
    temperature: float,
    max_tokens: int,
) -> tuple[str, str, str]:
    if not api_key.strip():
        return "Paste a VeloKey API key first.", "", ""
    if not model.strip():
        return "Enter a model ID available to your VeloKey account.", "", ""
    if not user_prompt.strip():
        return "Enter a user prompt.", "", ""

    messages: list[dict[str, str]] = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    messages.append({"role": "user", "content": user_prompt.strip()})

    payload = {
        "model": model.strip(),
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }

    try:
        status, body, elapsed_ms = _request_json("POST", CHAT_COMPLETIONS_URL, api_key, json=payload)
    except requests.RequestException as exc:
        return f"Request failed: {exc}", "", _format_json(payload)

    if status >= 400:
        return f"Chat completion returned HTTP {status}.", _format_json(body), _format_json(payload)

    answer = ""
    if isinstance(body, dict):
        choices = body.get("choices")
        if isinstance(choices, list) and choices:
            first = choices[0]
            if isinstance(first, dict):
                message = first.get("message")
                if isinstance(message, dict):
                    content = message.get("content")
                    if isinstance(content, str):
                        answer = content
                if not answer and isinstance(first.get("text"), str):
                    answer = str(first["text"])

    if not answer:
        answer = "Request succeeded, but the response format did not include choices[0].message.content."

    status_line = f"HTTP {status} in {elapsed_ms:.0f} ms"
    return status_line, answer, _format_json(body)


with gr.Blocks(
    title="VeloKey OpenAI-Compatible API Playground",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="green"),
    css="""
    .resource-links a { margin-right: 0.75rem; }
    .hint { color: #4b5563; font-size: 0.95rem; }
    """,
) as demo:
    gr.Markdown(
        """
        # VeloKey OpenAI-Compatible API Playground

        Test a VeloKey chat completion request from Hugging Face using your own API key.

        <p class="resource-links">
          <a href="https://velokey.ai?ref=huggingface-space" target="_blank">Website</a>
          <a href="https://docs.velokey.ai/api/introduction" target="_blank">API docs</a>
          <a href="https://velokey.ai/model?ref=huggingface-space" target="_blank">Models</a>
          <a href="https://velokey.ai/pricing?ref=huggingface-space" target="_blank">Pricing</a>
          <a href="https://velokey.ai/console/keys?ref=huggingface-space" target="_blank">Get API key</a>
        </p>
        """
    )

    with gr.Row():
        api_key_input = gr.Textbox(
            label="VeloKey API key",
            type="password",
            placeholder="vk-...",
            scale=2,
        )
        model_input = gr.Dropdown(
            label="Model ID",
            choices=MODEL_EXAMPLES,
            value=MODEL_EXAMPLES[0],
            allow_custom_value=True,
            scale=2,
        )

    with gr.Row():
        list_models_button = gr.Button("List available models", variant="secondary")
        model_status = gr.Textbox(label="Model list status", interactive=False)

    available_models = gr.Textbox(
        label="Available model IDs",
        lines=8,
        interactive=False,
        placeholder="Click List available models to query GET /v1/models.",
    )

    with gr.Accordion("Prompt settings", open=True):
        system_prompt_input = gr.Textbox(
            label="System prompt",
            value=DEFAULT_SYSTEM_PROMPT,
            lines=2,
        )
        user_prompt_input = gr.Textbox(
            label="User prompt",
            value=DEFAULT_USER_PROMPT,
            lines=5,
        )
        with gr.Row():
            temperature_input = gr.Slider(
                label="Temperature",
                minimum=0,
                maximum=2,
                step=0.1,
                value=0.7,
            )
            max_tokens_input = gr.Slider(
                label="Max tokens",
                minimum=16,
                maximum=2048,
                step=16,
                value=512,
            )

    run_button = gr.Button("Run chat completion", variant="primary")

    with gr.Row():
        status_output = gr.Textbox(label="Status", interactive=False)
        answer_output = gr.Textbox(label="Assistant response", lines=10, interactive=False)

    raw_json_output = gr.Code(label="Raw JSON response", language="json", lines=18)

    gr.Markdown(
        """
        <p class="hint">
        API keys are submitted with each request and are not saved by this app. For production apps,
        keep your VeloKey API key on your own backend, not in browser-side code.
        </p>
        """
    )

    list_models_button.click(
        fn=list_models,
        inputs=[api_key_input],
        outputs=[model_status, available_models],
    )

    run_button.click(
        fn=run_chat_completion,
        inputs=[
            api_key_input,
            model_input,
            system_prompt_input,
            user_prompt_input,
            temperature_input,
            max_tokens_input,
        ],
        outputs=[status_output, answer_output, raw_json_output],
    )


if __name__ == "__main__":
    demo.launch()