import os
import json
import gradio as gr
from openai import OpenAI
from datetime import datetime

ENV_API_KEY = os.getenv("LLMAPI_KEY")  # backup if user doesn't provide
DEFAULT_ENDPOINT = "https://api.llmapi.ai/v1"
DEFAULT_MODEL = "deepseek-v4-flash-0731"

def run_llm(endpoint: str, api_key: str, model: str, prompt: str):
    endpoint = (endpoint or "").strip().rstrip("/")
    api_key = (api_key or "").strip()
    model = (model or "").strip()
    prompt = (prompt or "").strip()

    if not endpoint:
        return ("Missing endpoint (e.g. https://api.llmapi.ai/v1).", None)

    key = api_key or (ENV_API_KEY or "").strip()
    if not key:
        return ("Missing API key. Provide one in the form or set Spaces secret env LLMAPI_KEY.", None)

    if not model:
        return ("Missing model (e.g. deepseek-v4-flash-0731).", None)

    if not prompt:
        return ("Missing prompt.", None)

    client = OpenAI(base_url=endpoint, api_key=key)

    try:
        resp = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
        )
        content = resp.choices[0].message.content

        # Build a JSON download payload (complete-ish)
        payload = {
            "created_at_utc": datetime.utcnow().isoformat() + "Z",
            "request": {
                "base_url": endpoint,
                "model": model,
                "messages": [{"role": "user", "content": prompt}],
            },
            "response": resp.model_dump() if hasattr(resp, "model_dump") else str(resp),
        }

        # Write to a file for Gradio download
        filename = "result.json"
        with open(filename, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)

        return (content, filename)
    except Exception as e:
        return (f"Error: {e}", None)

with gr.Blocks() as demo:
    gr.Markdown("# Custom LLM Endpoint (user-provided)")

    with gr.Row():
        endpoint = gr.Textbox(label="Endpoint (base_url)", value=DEFAULT_ENDPOINT)
        api_key = gr.Textbox(
            label="API Key (optional - uses env LLMAPI_KEY as backup)",
            type="password",
            placeholder="Leave blank to use env LLMAPI_KEY",
        )

    model = gr.Textbox(label="Model", value=DEFAULT_MODEL)
    prompt = gr.Textbox(label="Prompt", lines=6, placeholder="Type your prompt here...")

    run_btn = gr.Button("Send")

    out_text = gr.Textbox(label="Response", lines=10)
    out_file = gr.File(label="Download results as JSON")

    run_btn.click(
        fn=run_llm,
        inputs=[endpoint, api_key, model, prompt],
        outputs=[out_text, out_file],
    )

demo.launch()