import os

import gradio as gr

try:
    import spaces
except Exception:
    spaces = None

from model_adapters import create_adapter


MODEL_LABEL = os.getenv("RFAB_HISTORIC_MODEL_LABEL", os.getenv("RFAB_HISTORIC_MODEL_ID", "RFAB Historic Model"))
GPU_DURATION_SECONDS = int(os.getenv("RFAB_GPU_DURATION_SECONDS", "60"))

adapter = create_adapter()


def gpu(fn):
    if spaces is None:
        return fn
    return spaces.GPU(duration=GPU_DURATION_SECONDS)(fn)


def normalize_history(history):
    if not isinstance(history, list):
        return []

    normalized = []
    for message in history:
        if not isinstance(message, dict):
            continue
        role = message.get("role")
        if role not in {"user", "assistant"}:
            continue
        content = message.get("content") or []
        if isinstance(content, str):
            content = [{"type": "text", "text": content}]
        normalized.append({
            "role": role,
            "metadata": message.get("metadata"),
            "content": content,
            "options": message.get("options")
        })
    return normalized


@gpu
def _bot_reply(history, system_prompt="", temperature=0.7, max_tokens=256, top_p=1.0, top_k=0):
    history = normalize_history(history)
    answer = adapter.generate(
        history=history,
        system_prompt=system_prompt or "",
        temperature=float(temperature),
        max_tokens=int(max_tokens),
        top_p=float(top_p),
        top_k=int(top_k),
    )
    history.append({
        "role": "assistant",
        "metadata": None,
        "content": [{"type": "text", "text": answer}],
        "options": None
    })
    return history


def preview_reply(message, system_prompt, temperature, max_tokens):
    history = [{
        "role": "user",
        "metadata": None,
        "content": [{"type": "text", "text": message}],
        "options": None
    }]
    result = _bot_reply(history, system_prompt, temperature, max_tokens, 1.0, 0)
    return result[-1]["content"][0]["text"] if result else ""


with gr.Blocks(title=MODEL_LABEL) as demo:
    gr.Markdown(f"# {MODEL_LABEL}")
    gr.Markdown("Reality Fabricator Historic Chat Space. The backend uses the `_bot_reply` API endpoint.")

    with gr.Row():
        user_message = gr.Textbox(label="Message", value="What is electricity?")
    with gr.Row():
        system_prompt_box = gr.Textbox(label="System prompt", value="")
    with gr.Row():
        temperature_slider = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
        max_tokens_slider = gr.Slider(16, 1024, value=256, step=1, label="Max tokens")
    preview_button = gr.Button("Generate")
    preview_output = gr.Textbox(label="Response")

    preview_button.click(
        preview_reply,
        inputs=[user_message, system_prompt_box, temperature_slider, max_tokens_slider],
        outputs=preview_output,
        api_name=False,
    )

    history_input = gr.JSON(visible=False)
    system_prompt_input = gr.Textbox(visible=False)
    temperature_input = gr.Number(visible=False)
    max_tokens_input = gr.Number(visible=False)
    top_p_input = gr.Number(visible=False)
    top_k_input = gr.Number(visible=False)
    history_output = gr.JSON(visible=False)

    gr.Button("API", visible=False).click(
        _bot_reply,
        inputs=[
            history_input,
            system_prompt_input,
            temperature_input,
            max_tokens_input,
            top_p_input,
            top_k_input,
        ],
        outputs=history_output,
        api_name="_bot_reply",
    )


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