import random
from collections.abc import Mapping
from uuid import uuid4

from openai import OpenAI
import gradio as gr
import base64
import mimetypes
import copy
import os

# Workaround for PyCharm debugger + uvicorn compatibility error:
#   TypeError: _patch_asyncio.<locals>.run() got an unexpected keyword argument 'loop_factory'
DEBUG = False
if DEBUG is True:  # or sys.gettrace() is not None:  # Debugger is attached
    import asyncio
    _original_run = asyncio.run
    def _patched_run(main, **kwargs):
        kwargs.pop('loop_factory', None)  # Remove unsupported arg
        return _original_run(main, **kwargs)
    asyncio.run = _patched_run

from theme import apriel
from utils import COMMUNITY_POSTFIX_URL, get_model_config, check_format, models_config, \
    logged_event_handler, DEBUG_MODE, DEBUG_MODEL, log_debug, log_info, log_error, log_warning
from log_chat import log_chat

DEFAULT_MODEL_TEMPERATURE = 1.0
BUTTON_WIDTH = 160
DEFAULT_OPT_OUT_VALUE = DEBUG_MODE

# If DEBUG_MODEL is True, use an alternative model (without reasoning) for testing
# DEFAULT_MODEL_NAME = "Apriel-1.5-15B-thinker" if not DEBUG_MODEL else "Apriel-1.5-15B-thinker"
DEFAULT_MODEL_NAME = "Apriel-1.6-15B-Thinker"

SHOW_BANNER = False
INFO_BANNER_MARKDOWN = """
                    <span class="banner-message-text">ℹ️ This app has been updated to use the recommended temperature of 0.6. We had set it to 0.8 earlier and expect 0.6 to be better. Please provide feedback using the model link.</span>
                    """
NEW_MODEL_BANNER_MARKDOWN = """
                    <span class="banner-message-text"><span class="banner-message-emoji">🚀</span> Now running [Apriel-1.6-15B-Thinker](https://huggingface.co/ServiceNow-AI/Apriel-1.6-15b-Thinker) - 30% more efficient, frontier-class reasoning</span>
                    """
BANNER_MARKDOWN = NEW_MODEL_BANNER_MARKDOWN

BUTTON_ENABLED = gr.update(interactive=True)
BUTTON_DISABLED = gr.update(interactive=False)
INPUT_ENABLED = gr.update(interactive=True)
INPUT_DISABLED = gr.update(interactive=False)
DROPDOWN_ENABLED = gr.update(interactive=True)
DROPDOWN_DISABLED = gr.update(interactive=False)

SEND_BUTTON_ENABLED = gr.update(interactive=True, visible=True)
SEND_BUTTON_DISABLED = gr.update(interactive=True, visible=False)
STOP_BUTTON_ENABLED = gr.update(interactive=True, visible=True)
STOP_BUTTON_DISABLED = gr.update(interactive=True, visible=False)

chat_start_count = 0
model_config = {}
openai_client = None

USE_RANDOM_ENDPOINT = False
endpoint_rotation_count = 0

# Maximum number of image messages allowed per request
MAX_IMAGE_MESSAGES = 5


def app_loaded(state, request: gr.Request):
    message_html = setup_model(DEFAULT_MODEL_NAME, intial=False)
    state['session'] = request.session_hash if request else uuid4().hex
    log_debug(f"app_loaded() --> Session: {state['session']}")
    return state, message_html


def update_model_and_clear_chat(model_name):
    actual_model_name = model_name.replace("Model: ", "")
    desc = setup_model(actual_model_name)
    return desc, []


def setup_model(model_key, intial=False):
    global model_config, openai_client, endpoint_rotation_count
    model_config = get_model_config(model_key)
    log_debug(f"update_model() --> Model config: {model_config}")

    # ENVIRONMENT VARIABLE CONFIGURATION
    base_url = os.environ.get("API_BASE_URL")
    api_key = os.environ.get("API_KEY")

    if not base_url:
        raise ValueError("API_BASE_URL environment variable not set")
    if not api_key:
        raise ValueError("API_KEY environment variable not set")

    openai_client = OpenAI(
        api_key=api_key,
        base_url=base_url
    )
    model_config['base_url'] = base_url
    log_debug(f"Switched to model {model_key} using endpoint {base_url} (ENV VARS)")

    _model_hf_name = model_config.get("MODEL_HF_URL").split('https://huggingface.co/')[1]
    _link = f"<a href='{model_config.get('MODEL_HF_URL')}{COMMUNITY_POSTFIX_URL}' target='_blank'>{_model_hf_name}</a>"
    _description = f"We'd love to hear your thoughts on the model. Click here to provide feedback - {_link}"

    if intial:
        return
    else:
        return _description


def chat_started():
    # outputs: model_dropdown, user_input, send_btn, stop_btn, clear_btn, reasoning_effort
    return (DROPDOWN_DISABLED, gr.update(value="", interactive=False),
            SEND_BUTTON_DISABLED, STOP_BUTTON_ENABLED, BUTTON_DISABLED, gr.update(interactive=False))


def chat_finished():
    # outputs: model_dropdown, user_input, send_btn, stop_btn, clear_btn, reasoning_effort
    return DROPDOWN_ENABLED, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, gr.update(interactive=True)


def stop_chat(state):
    state["stop_flag"] = True
    gr.Info("Chat stopped")
    return state


def toggle_opt_out(state, checkbox):
    state["opt_out"] = checkbox
    return state


def run_chat_inference(history, message, state, reasoning_effort="medium"):
    global chat_start_count
    state["is_streaming"] = True
    state["stop_flag"] = False
    error = None
    # ENVIRONMENT VARIABLE MODEL
    model_name = os.environ.get("API_MODEL")
    if not model_name:
         raise ValueError("API_MODEL environment variable not set")
    # model_name = model_config.get('MODEL_NAME')
    temperature = model_config.get('TEMPERATURE', DEFAULT_MODEL_TEMPERATURE)
    output_tag_start = model_config.get('OUTPUT_TAG_START', "[BEGIN FINAL RESPONSE]")
    output_tag_end = model_config.get('OUTPUT_TAG_END', "[END FINAL RESPONSE]")
    output_stop_token = model_config.get('OUTPUT_STOP_TOKEN', "<|end|>")

    # Reinitialize the OpenAI client with a random endpoint from the list
    setup_model(model_config.get('MODEL_KEY'))
    log_info(f"Using model {model_name} (temperature: {temperature}, reasoning_effort: {reasoning_effort}) with endpoint {model_config.get('base_url')}")

    if len(history) == 0:
        state["chat_id"] = uuid4().hex

    if openai_client is None:
        log_info("Client UI is stale, letting user know to refresh the page")
        gr.Warning("Client UI is stale, please refresh the page")
        return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state

    # files will be the newly added files from the user
    files = []

    # outputs: model_dropdown, user_input, send_btn, stop_btn, clear_btn, session_state
    log_debug(f"{'-' * 80}")
    log_debug(f"chat_fn() --> Message: {message}")
    log_debug(f"chat_fn() --> History: {history}")

    # We have multimodal input in this case
    if isinstance(message, Mapping):
        files = message.get("files") or []
        message = message.get("text") or ""
        log_debug(f"chat_fn() --> Message (text only): {message}")
        log_debug(f"chat_fn() --> Files: {files}")

    # Validate that any uploaded files are images
    if len(files) > 0:
        invalid_files = []
        for path in files:
            try:
                mime, _ = mimetypes.guess_type(path)
                mime = mime or ""
                if not mime.startswith("image/"):
                    invalid_files.append((os.path.basename(path), mime or "unknown"))
            except Exception as e:
                log_error(f"Failed to inspect file '{path}': {e}")
                invalid_files.append((os.path.basename(path), "unknown"))

        if invalid_files:
            msg = "Only image files are allowed. Invalid uploads: " + \
                  ", ".join([f"{p} (type: {m})" for p, m in invalid_files])
            log_warning(msg)
            gr.Warning(msg)
            yield history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
            return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state

    # Enforce maximum number of files/images per request
    if len(files) > MAX_IMAGE_MESSAGES:
        gr.Warning(f"Too many images provided; keeping only the first {MAX_IMAGE_MESSAGES} file(s).")
        files = files[:MAX_IMAGE_MESSAGES]

    try:
        # Check if the message is empty
        if not message.strip() and len(files) == 0:
            gr.Info("Please enter a message before sending")
            yield history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
            return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state

        chat_start_count = chat_start_count + 1
        user_messages_count = sum(1 for item in history if isinstance(item, dict) and item.get("role") == "user"
                                  and isinstance(item.get("content"), str))
        log_info(f"chat_start_count: {chat_start_count}, turns: {user_messages_count + 1}, model: {model_name}")

        is_reasoning = model_config.get("REASONING")

        # Remove any assistant messages with metadata from history for multiple turns
        log_debug(f"Initial History: {history}")
        check_format(history, "messages")
        # Build UI history: add text (if any) and per-file image placeholders {"path": ...}
        # Build API parts separately later to avoid Gradio issues with arrays in content
        if len(files) == 0:
            history.append({"role": "user", "content": message})
        else:
            if message.strip():
                history.append({"role": "user", "content": message})
            for path in files:
                history.append({"role": "user", "content": {"path": path}})

        log_debug(f"History with user message: {history}")
        check_format(history, "messages")

        # Create the streaming response
        try:
            history_no_thoughts = [item for item in history if
                                   not (isinstance(item, dict) and
                                        item.get("role") == "assistant" and
                                        isinstance(item.get("metadata"), dict) and
                                        item.get("metadata", {}).get("title") is not None)]
            log_debug(f"Updated History: {history_no_thoughts}")
            check_format(history_no_thoughts, "messages")
            log_debug(f"history_no_thoughts with user message: {history_no_thoughts}")

            # Build API-specific messages:
            # - Convert any UI image placeholders {"path": ...} to image_url parts
            # - Convert any user string content that is a valid file path to image_url parts
            # - Coalesce consecutive image paths into a single image-only user message
            api_messages = []
            image_parts_buffer = []

            def flush_image_buffer():
                if len(image_parts_buffer) > 0:
                    api_messages.append({"role": "user", "content": list(image_parts_buffer)})
                    image_parts_buffer.clear()

            def to_image_part(path: str):
                try:
                    mime, _ = mimetypes.guess_type(path)
                    mime = mime or "application/octet-stream"
                    with open(path, "rb") as f:
                        b64 = base64.b64encode(f.read()).decode("utf-8")
                    data_url = f"data:{mime};base64,{b64}"
                    return {"type": "image_url", "image_url": {"url": data_url}}
                except Exception as e:
                    log_error(f"Failed to load file '{path}': {e}")
                    return None

            def normalize_msg(msg):
                # Returns (role, content, as_dict) where as_dict is a message dict suitable to pass through when unmodified
                if isinstance(msg, dict):
                    return msg.get("role"), msg.get("content"), msg
                # Gradio ChatMessage-like object
                role = getattr(msg, "role", None)
                content = getattr(msg, "content", None)
                if role is not None:
                    return role, content, {"role": role, "content": content}
                return None, None, msg

            for m in copy.deepcopy(history_no_thoughts):
                role, content, as_dict = normalize_msg(m)
                # Unknown structure: pass through
                if role is None:
                    flush_image_buffer()
                    api_messages.append(as_dict)
                    continue

                # Assistant messages pass through as-is
                if role == "assistant":
                    flush_image_buffer()
                    api_messages.append(as_dict)
                    continue

                # Only user messages have potential image paths to convert
                if role == "user":
                    # Case A: {'path': ...}
                    if isinstance(content, dict) and isinstance(content.get("path"), str):
                        p = content["path"]
                        part = to_image_part(p) if os.path.isfile(p) else None
                        if part:
                            image_parts_buffer.append(part)
                        else:
                            flush_image_buffer()
                            api_messages.append({"role": "user", "content": str(content)})
                        continue

                    # Case B: string or tuple content that may be a file path
                    if isinstance(content, str):
                        if os.path.isfile(content):
                            part = to_image_part(content)
                            if part:
                                image_parts_buffer.append(part)
                                continue
                        # Not a file path: pass through as text
                        flush_image_buffer()
                        api_messages.append({"role": "user", "content": content})
                        continue
                    if isinstance(content, tuple):
                        # Common case: a single-element tuple containing a path string
                        tuple_items = list(content)
                        tmp_parts = []
                        text_accum = []
                        for item in tuple_items:
                            if isinstance(item, str) and os.path.isfile(item):
                                part = to_image_part(item)
                                if part:
                                    tmp_parts.append(part)
                                else:
                                    text_accum.append(item)
                            else:
                                text_accum.append(str(item))
                        if tmp_parts:
                            flush_image_buffer()
                            api_messages.append({"role": "user", "content": tmp_parts})
                            if not text_accum:
                                continue
                        if text_accum:
                            flush_image_buffer()
                            api_messages.append({"role": "user", "content": "\n".join(text_accum)})
                            continue

                    # Case C: list content
                    if isinstance(content, list):
                        # If it's already a list of parts, let it pass through
                        all_dicts = all(isinstance(c, dict) for c in content)
                        if all_dicts:
                            flush_image_buffer()
                            api_messages.append({"role": "user", "content": content})
                            continue
                        # It might be a list of strings (paths/text). Convert string paths to image parts, others to text parts
                        tmp_parts = []
                        text_accum = []

                        def flush_text_accum():
                            if text_accum:
                                api_messages.append({"role": "user", "content": "\n".join(text_accum)})
                                text_accum.clear()
                        for item in content:
                            if isinstance(item, str) and os.path.isfile(item):
                                part = to_image_part(item)
                                if part:
                                    tmp_parts.append(part)
                                else:
                                    text_accum.append(item)
                            else:
                                text_accum.append(str(item))
                        if tmp_parts:
                            flush_image_buffer()
                            api_messages.append({"role": "user", "content": tmp_parts})
                        if text_accum:
                            flush_text_accum()
                        continue

                    # Fallback: pass through
                    flush_image_buffer()
                    api_messages.append(as_dict)
                    continue

                # Other roles
                flush_image_buffer()
                api_messages.append(as_dict)

            # Flush any trailing images
            flush_image_buffer()

            log_debug(f"sending api_messages to model {model_name}: {api_messages}")

            # Ensure we don't send too many images (count only messages whose content is a list of parts)
            image_msg_indices = [
                i for i, msg in enumerate(api_messages)
                if isinstance(msg, dict) and isinstance(msg.get('content'), list)
            ]
            image_count = len(image_msg_indices)
            if image_count > MAX_IMAGE_MESSAGES:
                # Remove oldest image messages until we have MAX_IMAGE_MESSAGES or fewer
                to_remove = image_count - MAX_IMAGE_MESSAGES
                removed = 0
                for idx in image_msg_indices:
                    if removed >= to_remove:
                        break
                    # Pop considering prior removals shift indices
                    api_messages.pop(idx - removed)
                    removed += 1
                gr.Warning(f"Too many images provided; keeping the latest {MAX_IMAGE_MESSAGES} and dropped {removed} older image message(s).")

            stream = openai_client.chat.completions.create(
                model=model_name,
                messages=api_messages,
                temperature=temperature,
                top_p=1.0,
                reasoning_effort=reasoning_effort,
                stream=True
            )
        except Exception as e:
            log_error(f"Error:\n\t{e}\n\tInference failed for model {model_name} and endpoint {model_config['base_url']}")
            error = str(e)
            yield ([{"role": "assistant",
                     "content": "😔 The model is unavailable at the moment. Please try again later."}],
                   INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state)
            if state["opt_out"] is not True:
                log_chat(chat_id=state["chat_id"],
                         session_id=state["session"],
                         model_name=model_name,
                         prompt=message,
                         history=history,
                         info={"is_reasoning": model_config.get("REASONING"), "temperature": temperature,
                               "stopped": True, "error": str(e)},
                         )
            else:
                log_info(f"User opted out of chat history. Not logging chat. model: {model_name}")
            return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state

        if is_reasoning:
            history.append(gr.ChatMessage(
                role="assistant",
                content="Thinking...",
                metadata={"title": "🧠 Thought"}
            ))
            log_debug(f"History added thinking: {history}")
            check_format(history, "messages")
        else:
            history.append(gr.ChatMessage(
                role="assistant",
                content="",
            ))
            log_debug(f"History added empty assistant: {history}")
            check_format(history, "messages")

        output_reasoning = ""
        output_content = ""
        completion_started = False
        
        for chunk in stream:
            if state["stop_flag"]:
                log_debug(f"chat_fn() --> Stopping streaming...")
                break  # Exit the loop if the stop flag is set
            
            delta = chunk.choices[0].delta
            new_reasoning = getattr(delta, "reasoning_content", "") or ""
            new_content = getattr(delta, "content", "") or ""
            
            output_reasoning += new_reasoning
            output_content += new_content

            if is_reasoning:
                # Update the reasoning bubble
                history[-1 if not completion_started else -2] = gr.ChatMessage(
                    role="assistant",
                    content=output_reasoning,
                    metadata={"title": "🧠 Thought"}
                )
                
                # Handle the content bubble
                # Check if we have actual content or if we should start the content bubble
                if new_content or (output_content and not completion_started):
                    # Clean up stop tokens from the content if present
                    if output_tag_end and output_content.endswith(output_tag_end):
                        output_content = output_content.replace(output_tag_end, "")
                    if output_stop_token and output_content.endswith(output_stop_token):
                        output_content = output_content.replace(output_stop_token, "")
                    
                    if not completion_started:
                        completion_started = True
                        history.append(gr.ChatMessage(
                            role="assistant",
                            content=output_content
                        ))
                    else:
                        history[-1] = gr.ChatMessage(
                            role="assistant",
                            content=output_content
                        )
            else:
                if output_content.endswith("<|end|>"):
                    output_content = output_content.replace("<|end|>", "")
                if output_content.endswith("<|end|>\n"):
                    output_content = output_content.replace("<|end|>\n", "")
                history[-1] = gr.ChatMessage(
                    role="assistant",
                    content=output_content
                )

            # log_message(f"Yielding messages: {history}")
            yield history, INPUT_DISABLED, SEND_BUTTON_DISABLED, STOP_BUTTON_ENABLED, BUTTON_DISABLED, state

        log_debug(f"Final History: {history}")
        check_format(history, "messages")
        yield history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state
    finally:
        if error is None:
            log_debug(f"chat_fn() --> Finished streaming. {chat_start_count} chats started.")
            if state["opt_out"] is not True:
                log_chat(chat_id=state["chat_id"],
                         session_id=state["session"],
                         model_name=model_name,
                         prompt=message,
                         history=history,
                         info={"is_reasoning": model_config.get("REASONING"), "temperature": temperature,
                               "stopped": state["stop_flag"]},
                         )

            else:
                log_info(f"User opted out of chat history. Not logging chat. model: {model_name}")
        state["is_streaming"] = False
        state["stop_flag"] = False
        return history, INPUT_ENABLED, SEND_BUTTON_ENABLED, STOP_BUTTON_DISABLED, BUTTON_ENABLED, state


log_info(f"Gradio version: {gr.__version__}")

title = None
description = None
theme = apriel

with open('styles.css', 'r') as f:
    custom_css = f.read()

with gr.Blocks(theme=theme, css=custom_css) as demo:
    session_state = gr.State(value={
        "is_streaming": False,
        "stop_flag": False,
        "chat_id": None,
        "session": None,
        "opt_out": DEFAULT_OPT_OUT_VALUE,
        "agreed": False,
    })  # Store session state as a dictionary

    gr.HTML(f"""
        <style>
            @media (min-width: 1024px) {{
                .send-button-container, .clear-button-container {{
                    max-width: {BUTTON_WIDTH}px;
                }}
            }}
        </style>
        """, elem_classes="css-styles")
    if SHOW_BANNER:
        with gr.Row(variant="compact", elem_classes=["responsive-row", "no-padding"], ):
            with gr.Column():
                gr.Markdown(BANNER_MARKDOWN, elem_classes="banner-message")

    with gr.Row(variant="panel", elem_classes="responsive-row", visible=False):
        with gr.Column(scale=1, min_width=400, elem_classes="model-dropdown-container"):
            model_dropdown = gr.Dropdown(
                choices=[f"Model: {model}" for model in models_config.keys()],
                value=f"Model: {DEFAULT_MODEL_NAME}",
                label=None,
                interactive=True,
                container=False,
                scale=0,
                min_width=400
            )
        with gr.Column(scale=4, min_width=0):
            feedback_message_html = gr.HTML(description, elem_classes="model-message")

    with gr.Column(visible=True, elem_classes="agreement-overlay") as agreement_overlay:
        with gr.Column(elem_classes="form"):
            gr.Markdown("## Privacy Agreement")
            gr.Markdown("""
By using this app, you agree to the following terms:

We record all content you submit and all model outputs (“Data”), including text, images, files, and minimal request metadata (timestamp & technical logs). We do not store IP addresses, cookies, or account identifiers, so we cannot link any submission back to a particular person. However, the text you submit may itself contain personal information (e.g., names, Social Security numbers). Please do not include sensitive personal data in your prompts. Any such information will be subject to our redaction process before any public release.

Data is used for research, safety evaluation, and to improve the Service. We reserve the right to publish, share, or redistribute redacted versions of the Data under a Creative Commons Attribution (CC‑BY) or similar open license. Before any public release, we apply automated and manual redaction to remove private keys, names, contact details, and other identifiers that may appear in the content.

Because we do not track user identities, individual submissions cannot be deleted or withdrawn once made. If you do not want your content used or released, do not submit it.
""")
            agree_btn = gr.Button("I Agree", variant="primary")

    with gr.Column(visible=True) as main_app_area:
        chatbot = gr.Chatbot(
            type="messages",
            height="calc(100svh - 320px)",
            max_height="calc(100svh - 320px)",
            elem_classes="chatbot",
        )

        with gr.Row():
            with gr.Column(scale=10, min_width=400, elem_classes="user-input-container"):
                with gr.Row():
                    # user_input = gr.MultimodalTextbox(
                    #     interactive=True,
                    #     container=False,
                    #     file_count="multiple",
                    #     placeholder="Type your message here and press Enter or upload file...",
                    #     show_label=False,
                    #     sources=["upload"],
                    #     max_plain_text_length=100000,
                    #     max_lines=10
                    # )

                    # Original text-only input
                    user_input = gr.Textbox(
                        show_label=False,
                        placeholder="Type your message here and press Enter",
                        container=False,
                        max_lines=10
                    )
            with gr.Column(scale=1, min_width=BUTTON_WIDTH * 2 + 20):
                with gr.Row():
                    with gr.Column(scale=1, min_width=BUTTON_WIDTH, elem_classes="send-button-container"):
                        send_btn = gr.Button("Send", variant="primary", elem_classes="control-button")
                        stop_btn = gr.Button("Stop", variant="cancel", elem_classes="control-button", visible=False)
                    with gr.Column(scale=1, min_width=BUTTON_WIDTH, elem_classes="clear-button-container"):
                        clear_btn = gr.ClearButton(chatbot, value="New Chat", variant="secondary", elem_classes="control-button")
        with gr.Row():
             with gr.Column(scale=1):
                  reasoning_effort_radio = gr.Radio(
                      choices=["low", "medium", "high"],
                      value="medium",
                      label="Reasoning Effort",
                      interactive=True,
                      container=True,
                      elem_classes="reasoning-radio"
                  )


    def agree_to_terms(state):
        log_info("Privacy agreement accepted by user")
        state["agreed"] = True
        return gr.update(visible=False), state

    # Use JavaScript to directly hide the overlay - bypasses Gradio's state management
    # which can be unreliable on HuggingFace Spaces
    agree_btn.click(
        agree_to_terms,
        inputs=[session_state],
        outputs=[agreement_overlay, session_state],
        queue=False,
        js="() => { document.querySelector('.agreement-overlay').style.display = 'none'; }"
    )

    gr.on(
        triggers=[send_btn.click, user_input.submit],
        fn=run_chat_inference,  # this generator streams results. do not use logged_event_handler wrapper
        inputs=[chatbot, user_input, session_state, reasoning_effort_radio],
        outputs=[chatbot, user_input, send_btn, stop_btn, clear_btn, session_state],
        concurrency_limit=4,
        api_name=False
    ).then(
        fn=chat_finished, inputs=None, outputs=[model_dropdown, user_input, send_btn, stop_btn, clear_btn, reasoning_effort_radio], queue=False)

    # In parallel, disable or update the UI controls
    gr.on(
        triggers=[send_btn.click, user_input.submit],
        fn=chat_started,
        inputs=None,
        outputs=[model_dropdown, user_input, send_btn, stop_btn, clear_btn, reasoning_effort_radio],
        queue=False,
        show_progress='hidden',
        api_name=False
    )

    stop_btn.click(
        fn=stop_chat,
        inputs=[session_state],
        outputs=[session_state],
        api_name=False
    )



    # Ensure the model is reset to default on page reload
    demo.load(
        fn=logged_event_handler(
            log_msg="Browser session started",
            event_handler=app_loaded
        ),
        inputs=[session_state],
        outputs=[session_state, feedback_message_html],
        queue=True,
        api_name=False
    )

    model_dropdown.change(
        fn=update_model_and_clear_chat,
        inputs=[model_dropdown],
        outputs=[feedback_message_html, chatbot],
        api_name=False
    )

demo.queue(default_concurrency_limit=2).launch(ssr_mode=False, show_api=False, max_file_size="10mb")
log_info("Gradio app launched")