import gradio as gr
from azure.ai.inference import ChatCompletionsClient
from azure.ai.inference.models import SystemMessage, UserMessage, AssistantMessage
from azure.core.credentials import AzureKeyCredential


def _normalize_endpoint(endpoint: str) -> str:
    """
    ChatCompletionsClient expects the Azure AI inference endpoint that ends with /models.
    Many UIs show a "project endpoint" or a resource endpoint without /models.
    This function tries to normalize what the user pastes into a usable inference endpoint.
    """
    endpoint = (endpoint or "").strip()
    if not endpoint:
        raise gr.Error("Please provide an Azure/Microsoft Foundry endpoint.")

    if not endpoint.startswith("http"):
        endpoint = "https://" + endpoint

    endpoint = endpoint.rstrip("/")

    # If it's an Azure AI Foundry / Azure AI Services inference host, ensure /models is present
    if ".services.ai.azure.com" in endpoint and not endpoint.endswith("/models"):
        endpoint = endpoint + "/models"

    return endpoint


def _build_client(endpoint: str, api_key: str) -> ChatCompletionsClient:
    api_key = (api_key or "").strip()
    if not api_key:
        raise gr.Error("Please provide an API key.")

    endpoint = _normalize_endpoint(endpoint)
    return ChatCompletionsClient(endpoint=endpoint, credential=AzureKeyCredential(api_key))


def _to_messages(system_prompt: str, history: list[tuple[str, str]], user_text: str):
    msgs = []
    system_prompt = (system_prompt or "").strip()
    if system_prompt:
        msgs.append(SystemMessage(content=system_prompt))

    for u, a in (history or []):
        if u:
            msgs.append(UserMessage(content=u))
        if a:
            msgs.append(AssistantMessage(content=a))

    msgs.append(UserMessage(content=(user_text or "")))
    return msgs


def chat(
    endpoint: str,
    api_key: str,
    model: str,
    system_prompt: str,
    temperature: float,
    max_tokens: int,
    top_p: float,
    user_text: str,
    history: list[tuple[str, str]],
):
    user_text = (user_text or "").strip()
    if not user_text:
        return "", history, ""

    # Build client (normalizes endpoint and checks key)
    client = _build_client(endpoint, api_key)

    # Important: In many Foundry setups, "model" must be the DEPLOYMENT NAME,
    # not the underlying base model string.
    model = (model or "").strip()
    if not model:
        raise gr.Error("Please provide a model/deployment name (e.g., your deployment for gpt-4o).")

    try:
        resp = client.complete(
            model=model,
            messages=_to_messages(system_prompt, history, user_text),
            temperature=float(temperature),
            max_tokens=int(max_tokens),
            top_p=float(top_p),
        )
        answer = resp.choices[0].message.content or ""
        used_endpoint = client._endpoint  # for display/debug in the UI
        info = f"✅ Called: {used_endpoint}\n✅ Model/Deployment: {model}"
    except Exception as e:
        raise gr.Error(
            "Azure call failed.\n\n"
            "Most common fixes:\n"
            "1) Ensure your endpoint is the Azure AI inference endpoint (it should end with /models).\n"
            "2) Ensure the key matches that same resource/project.\n"
            "3) Ensure 'Model/Deployment' is your DEPLOYMENT NAME in Foundry.\n\n"
            f"Error: {type(e).__name__}: {e}"
        )

    history = (history or []) + [(user_text, answer)]
    return "", history, info


def clear_all():
    return [], ""


with gr.Blocks(title="Azure (Foundry) GPT-4o Chatbot") as demo:
    gr.Markdown(
        "## Azure / Microsoft Foundry Chatbot\n"
        "**Tip:** Your endpoint must be the Azure AI inference endpoint. If you paste a Foundry resource endpoint like\n"
        "`https://xxxx.services.ai.azure.com`, this app will automatically append `/models`.\n\n"
        "**Model/Deployment:** often must be your **deployment name** (even if the base model is GPT-4o)."
    )

    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(height=520)
            user_text = gr.Textbox(
                label="Message",
                placeholder="Type your message and press Enter…",
                lines=2,
            )
            with gr.Row():
                send = gr.Button("Send", variant="primary")
                clear = gr.Button("Clear chat")

        with gr.Column(scale=3):
            with gr.Accordion("Connection (Endpoint + Key)", open=True):
                endpoint = gr.Textbox(
                    label="Azure endpoint (Foundry resource or inference endpoint)",
                    placeholder="https://projectxxxxx-resource.services.ai.azure.com  (app will add /models)",
                )
                api_key = gr.Textbox(
                    label="API key",
                    placeholder="Paste your key here",
                    type="password",
                )

            with gr.Accordion("Model + generation settings", open=True):
                model = gr.Textbox(
                    label="Model / Deployment name",
                    value="gpt-4o",
                    info="If you get 'model not found', change this to your deployment name from Foundry."
                )
                system_prompt = gr.Textbox(
                    label="System prompt (optional)",
                    value="You are a helpful assistant.",
                    lines=3,
                )
                temperature = gr.Slider(0, 1.5, value=0.7, step=0.1, label="Temperature")
                max_tokens = gr.Slider(64, 4096, value=1024, step=64, label="Max tokens")
                top_p = gr.Slider(0.1, 1.0, value=1.0, step=0.05, label="Top-p")

            debug_info = gr.Textbox(
                label="Debug info",
                value="",
                interactive=False,
                lines=3,
            )

    send.click(
        chat,
        inputs=[endpoint, api_key, model, system_prompt, temperature, max_tokens, top_p, user_text, chatbot],
        outputs=[user_text, chatbot, debug_info],
    )
    user_text.submit(
        chat,
        inputs=[endpoint, api_key, model, system_prompt, temperature, max_tokens, top_p, user_text, chatbot],
        outputs=[user_text, chatbot, debug_info],
    )

    clear.click(clear_all, outputs=[chatbot, debug_info])

demo.launch(server_name="0.0.0.0", server_port=7860)