import os
import requests
import threading
from typing import Optional, List, Tuple

import gradio as gr


ENDPOINT_URL = "https://austrian-code-wizard--metarlaif-web-dev.modal.run"


def get_feedback_options() -> List[Tuple[str, str]]:
    args = {
        "C3PO_API_KEY": os.environ.get("C3PO_API_KEY"),
    }
    response = requests.post(f"{ENDPOINT_URL}/list_adapters", json=args)
    data = response.json()["adapters"]
    return [
        (adapter["feedback_name"], adapter["feedback_id"])
    for adapter in data]


def get_completion(prompt: str, adapter: Optional[str] = None) -> str:
    args = {
        "C3PO_API_KEY": os.environ.get("C3PO_API_KEY"),
        "prompt": prompt,
        "adapter": adapter,
    }
    response = requests.post(f"{ENDPOINT_URL}/completion", json=args)
    data = response.json()
    return data["response"]


def warmup(*args):
    args = {
        "C3PO_API_KEY": os.environ.get("C3PO_API_KEY"),
    }
    # Warmup the server but don't wait for the response
    threading.Thread(target=requests.post, args=(f"{ENDPOINT_URL}/warmup"), kwargs={"json": args}, daemon=True).start()

dropdown_options = get_feedback_options()

demo = gr.Interface(
    get_completion,
    [
        gr.Textbox(
            placeholder="Enter a prompt...", label="Prompt"
        ),
        gr.Dropdown(
            choices=dropdown_options, label="Feedback", info="Will add the adapter for the respective feedback to the model."
        )
    ],
    "text",
    concurrency_limit=8
)

if __name__ == "__main__":
    demo.queue(max_size=32)
    demo.launch()