import json
import os
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr
import modelscope_studio.components.antd as antd
import modelscope_studio.components.antdx as antdx
import modelscope_studio.components.base as ms
import modelscope_studio.components.pro as pro
from huggingface_hub import CommitScheduler
from openai import OpenAI

from config import (
    CHAR_YIELD_SIZE,
    DEFAULT_LOCALE,
    DEFAULT_SETTINGS,
    DEFAULT_THEME,
    KEPT_TURNS,
    MODEL_OPTIONS_MAP,
    bot_config,
    save_history,
    user_config,
    welcome_config,
)
from interface.logo import Logo
from interface.settings import SettingsMenu

logs_id = os.getenv("LOGS_ID")
logs_token = os.getenv("HF_LOGS_TOKEN")

logs_file = Path("logs/") / f"data_{uuid.uuid4()}.json"
logs_folder = logs_file.parent

scheduler = CommitScheduler(
    repo_id=logs_id,
    repo_type="dataset",
    folder_path=logs_folder,
    path_in_repo="data",
    every=5,
    token=logs_token,
    private=True,
)


def format_history(history):
    messages = []
    for item in history:
        if item["role"] == "user":
            messages.append({"role": "user", "content": item["content"]})
        elif item["role"] == "assistant":
            contents = [
                {"type": "text", "text": content["content"]}
                for content in item["content"]
                if content["type"] == "text"
            ]
            messages.append(
                {
                    "role": "assistant",
                    "content": contents[0]["text"] if len(contents) > 0 else "",
                }
            )
    return messages


class Gradio_Events:
    @staticmethod
    def submit(
        state_value,
        request: gr.Request,
        profile: gr.OAuthProfile | None,
    ):
        settings = state_value["conversation_contexts"][state_value["conversation_id"]][
            "settings"
        ]

        model = settings.get("model")
        api_key_env = MODEL_OPTIONS_MAP.get(model).get("api_key_env")
        base_url_env = MODEL_OPTIONS_MAP.get(model).get("base_url_env")
        api_key_value = os.getenv(api_key_env)
        base_url_value = os.getenv(base_url_env)

        client = OpenAI(api_key=api_key_value, base_url=base_url_value)

        history = state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ]
        messages = format_history(history)

        history.append(
            {
                "role": "assistant",
                "content": [],
                "key": str(uuid.uuid4()),
                "header": MODEL_OPTIONS_MAP.get(model, {}).get("label", None),
                "loading": True,
                "status": "pending",
            }
        )

        yield {
            chatbot: gr.update(value=history),
            state: gr.update(value=state_value),
        }

        try:
            kept_messages = messages[:2] + messages[-(2 * (KEPT_TURNS - 1) + 1) :]
            response = client.chat.completions.create(
                model=model,
                messages=kept_messages,
                temperature=settings.get("temperature"),
                top_p=settings.get("top_p"),
                max_tokens=settings.get("max_new_tokens"),
                stream=True,
            )
            start_time = time.time()
            thought_cost_time = 0
            reasoning_content = ""
            answer_content = ""
            is_thinking = False
            is_answering = False
            contents = [None, None]
            last_yielded_length = 0
            for chunk in response:
                if not getattr(
                    getattr(chunk.choices[0], "message", None), "content", None
                ) and not not getattr(
                    getattr(chunk.choices[0], "message", None),
                    "reasoning_content",
                    None,
                ):
                    pass

                else:
                    delta = chunk.choices[0].delta
                    if hasattr(delta, "reasoning_content") and delta.reasoning_content:
                        if not is_thinking:
                            contents[0] = {
                                "type": "tool",
                                "content": "",
                                "options": {
                                    "title": "Let me think...",
                                    "status": "pending",
                                },
                                "copyable": False,
                                "editable": False,
                            }
                            is_thinking = True
                        reasoning_content += delta.reasoning_content
                    if hasattr(delta, "content") and delta.content:
                        if not is_answering:
                            thought_cost_time = "{:.2f}".format(
                                time.time() - start_time
                            )
                            if contents[0]:
                                contents[0]["options"]["title"] = (
                                    f"Inner reflections ({thought_cost_time}s)"
                                )
                                contents[0]["options"]["status"] = "done"
                            contents[1] = {
                                "type": "text",
                                "content": "",
                            }
                            is_answering = True
                        answer_content += delta.content
                    if contents[0]:
                        contents[0]["content"] = reasoning_content
                    if contents[1]:
                        contents[1]["content"] = answer_content
                current_length = len(reasoning_content) + len(answer_content)
                if current_length - last_yielded_length >= CHAR_YIELD_SIZE:
                    last_yielded_length = current_length
                    history[-1]["content"] = [
                        content for content in contents if content
                    ]
                    history[-1]["loading"] = False
                    yield {
                        chatbot: gr.update(value=history),
                    }
            history[-1]["content"] = [content for content in contents if content]
            history[-1]["loading"] = False
            history[-1]["status"] = "done"
            cost_time = "{:.2f}".format(time.time() - start_time)
            history[-1]["footer"] = f"{cost_time}s"
            yield {
                chatbot: gr.update(value=history),
                state: gr.update(value=state_value),
            }
            with scheduler.lock:
                with logs_file.open("a") as f:
                    f.write(
                        json.dumps(
                            {
                                "request_id": str(uuid.uuid4()),
                                "input": messages,
                                "reasoning_content": reasoning_content,
                                "output": answer_content,
                                "turn": len(messages),
                                "model": model,
                                "temperature": settings.get("temperature"),
                                "top_p": settings.get("top_p"),
                                "max_tokens": settings.get("max_new_tokens"),
                                "latency": cost_time,
                                "reasoning_latency": thought_cost_time
                                if thought_cost_time
                                else cost_time,
                                "timestamp": datetime.now(timezone.utc).isoformat(),
                                "hf_username": profile.username if profile else None,
                                "session_hash": request.session_hash,
                                "client_ip": request.client.host,
                                "user_agent": request.headers.get("user-agent"),
                                "referer": request.headers.get("referer"),
                                "accept_language": request.headers.get(
                                    "accept-language"
                                ),
                            }
                        )
                    )
                    f.write("\n")
        except Exception as e:
            print("model: ", model, "-", "Error: ", e)
            history[-1]["loading"] = False
            history[-1]["status"] = "done"
            history[-1]["content"] += [
                {
                    "type": "text",
                    "content": f'<span style="color: var(--color-red-500)">{str(e)}</span>',
                }
            ]
            yield {
                chatbot: gr.update(value=history),
                state: gr.update(value=state_value),
            }
            raise e

    @staticmethod
    def add_message(
        input_value,
        settings_form_value,
        state_value,
        request: gr.Request,
        profile: gr.OAuthProfile | None,
    ):
        if not state_value["conversation_id"]:
            random_id = str(uuid.uuid4())
            history = []
            state_value["conversation_id"] = random_id
            state_value["conversation_contexts"][state_value["conversation_id"]] = {
                "history": history
            }
            state_value["conversations"].append(
                {"label": input_value, "key": random_id}
            )

        history = state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ]

        state_value["conversation_contexts"][state_value["conversation_id"]] = {
            "history": history,
            "settings": settings_form_value,
        }
        history.append(
            {"role": "user", "content": input_value, "key": str(uuid.uuid4())}
        )
        yield Gradio_Events.preprocess_submit(clear_input=True)(state_value)

        try:
            for chunk in Gradio_Events.submit(state_value, request, profile):
                yield chunk
        except Exception as e:
            raise e
        finally:
            yield Gradio_Events.postprocess_submit(state_value)

    @staticmethod
    def preprocess_submit(clear_input=True):
        def preprocess_submit_handler(state_value):
            history = state_value["conversation_contexts"][
                state_value["conversation_id"]
            ]["history"]
            return {
                **(
                    {
                        input: gr.update(value=None, loading=True)
                        if clear_input
                        else gr.update(loading=True),
                    }
                    if clear_input
                    else {}
                ),
                conversations: gr.update(
                    active_key=state_value["conversation_id"],
                    items=list(
                        map(
                            lambda item: {
                                **item,
                                "disabled": True
                                if item["key"] != state_value["conversation_id"]
                                else False,
                            },
                            state_value["conversations"],
                        )
                    ),
                ),
                add_conversation_btn: gr.update(disabled=True),
                clear_btn: gr.update(disabled=True),
                conversation_delete_menu_item: gr.update(disabled=True),
                chatbot: gr.update(
                    value=history,
                    bot_config=bot_config(disabled_actions=["edit", "retry", "delete"]),
                    user_config=user_config(disabled_actions=["edit", "delete"]),
                ),
                state: gr.update(value=state_value),
            }

        return preprocess_submit_handler

    @staticmethod
    def postprocess_submit(state_value):
        history = state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ]
        return {
            input: gr.update(loading=False),
            conversation_delete_menu_item: gr.update(disabled=False),
            clear_btn: gr.update(disabled=False),
            conversations: gr.update(items=state_value["conversations"]),
            add_conversation_btn: gr.update(disabled=False),
            chatbot: gr.update(
                value=history, bot_config=bot_config(), user_config=user_config()
            ),
            state: gr.update(value=state_value),
        }

    @staticmethod
    def cancel(state_value):
        history = state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ]
        history[-1]["loading"] = False
        history[-1]["status"] = "done"
        history[-1]["footer"] = "Chat completion paused"
        return Gradio_Events.postprocess_submit(state_value)

    @staticmethod
    def delete_message(state_value, e: gr.EventData):
        index = e._data["payload"][0]["index"]
        history = state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ]
        history = history[:index] + history[index + 1 :]

        state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ] = history

        return gr.update(value=state_value)

    @staticmethod
    def edit_message(state_value, chatbot_value, e: gr.EventData):
        index = e._data["payload"][0]["index"]
        history = state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ]
        history[index]["content"] = chatbot_value[index]["content"]
        return gr.update(value=state_value)

    @staticmethod
    def regenerate_message(
        settings_form_value,
        state_value,
        e: gr.EventData,
        request: gr.Request,
        profile: gr.OAuthProfile | None,
    ):
        index = e._data["payload"][0]["index"]
        history = state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ]
        history = history[:index]

        state_value["conversation_contexts"][state_value["conversation_id"]] = {
            "history": history,
            "settings": settings_form_value,
        }

        yield Gradio_Events.preprocess_submit()(state_value)
        try:
            for chunk in Gradio_Events.submit(state_value, request, profile):
                yield chunk
        except Exception as e:
            raise e
        finally:
            yield Gradio_Events.postprocess_submit(state_value)

    @staticmethod
    def select_suggestion(input_value, e: gr.EventData):
        input_value = input_value[:-1] + e._data["payload"][0]
        return gr.update(value=input_value)

    @staticmethod
    def apply_prompt(e: gr.EventData):
        return gr.update(value=e._data["payload"][0]["value"]["description"])

    @staticmethod
    def new_chat(state_value):
        if not state_value["conversation_id"]:
            return gr.skip()
        state_value["conversation_id"] = ""
        return (
            gr.update(active_key=state_value["conversation_id"]),
            gr.update(value=None),
            gr.update(value=DEFAULT_SETTINGS),
            gr.update(value=state_value),
        )

    @staticmethod
    def select_conversation(state_value, e: gr.EventData):
        active_key = e._data["payload"][0]
        if state_value["conversation_id"] == active_key or (
            active_key not in state_value["conversation_contexts"]
        ):
            return gr.skip()
        state_value["conversation_id"] = active_key
        return (
            gr.update(active_key=active_key),
            gr.update(
                value=state_value["conversation_contexts"][active_key]["history"]
            ),
            gr.update(
                value=state_value["conversation_contexts"][active_key]["settings"]
            ),
            gr.update(value=state_value),
        )

    @staticmethod
    def click_conversation_menu(state_value, e: gr.EventData):
        conversation_id = e._data["payload"][0]["key"]
        operation = e._data["payload"][1]["key"]
        if operation == "delete":
            del state_value["conversation_contexts"][conversation_id]

            state_value["conversations"] = [
                item
                for item in state_value["conversations"]
                if item["key"] != conversation_id
            ]

            if state_value["conversation_id"] == conversation_id:
                state_value["conversation_id"] = ""
                return (
                    gr.update(
                        items=state_value["conversations"],
                        active_key=state_value["conversation_id"],
                    ),
                    gr.update(value=None),
                    gr.update(value=state_value),
                )
            else:
                return (
                    gr.update(items=state_value["conversations"]),
                    gr.skip(),
                    gr.update(value=state_value),
                )
        return gr.skip()

    @staticmethod
    def toggle_settings_header(settings_header_state_value):
        settings_header_state_value["open"] = not settings_header_state_value["open"]
        return gr.update(value=settings_header_state_value)

    @staticmethod
    def clear_conversation_history(state_value):
        if not state_value["conversation_id"]:
            return gr.skip()
        state_value["conversation_contexts"][state_value["conversation_id"]][
            "history"
        ] = []
        return gr.update(value=None), gr.update(value=state_value)

    @staticmethod
    def update_browser_state(state_value):
        return gr.update(
            value=dict(
                conversations=state_value["conversations"],
                conversation_contexts=state_value["conversation_contexts"],
            )
        )

    @staticmethod
    def apply_browser_state(browser_state_value, state_value):
        state_value["conversations"] = browser_state_value["conversations"]
        state_value["conversation_contexts"] = browser_state_value[
            "conversation_contexts"
        ]
        return gr.update(items=browser_state_value["conversations"]), gr.update(
            value=state_value
        )


css = """
.gradio-container {
  padding: 0 !important;
}

.gradio-container > main.fillable {
  padding: 0 !important;
}

#chatbot {
  height: calc(100vh - 21px - 16px);
  max-height: 1500px;
}

#chatbot .chatbot-conversations {
  height: 100vh;
  background-color: var(--ms-gr-ant-color-bg-layout);
  padding-left: 4px;
  padding-right: 4px;
}

#chatbot .chatbot-conversations .chatbot-conversations-list {
  padding-left: 0;
  padding-right: 0;
}

/* NEW: Right sidebar styles */
#chatbot .chatbot-settings-sidebar {
  height: 100vh;
  background-color: var(--ms-gr-ant-color-bg-layout);
  overflow-y: auto;
}

#chatbot .chatbot-settings-sidebar .ant-card {
  height: 100%;
}

#chatbot .chatbot-chat {
  padding: 32px;
  padding-bottom: 0;
  height: 100%;
}

@media (max-width: 768px) {
  #chatbot .chatbot-chat {
      padding: 0;
  }
}

#chatbot .chatbot-chat .chatbot-chat-messages {
  flex: 1;
}

#chatbot .setting-form-temperature .ms-gr-ant-form-item-control-input-content {
    display: flex;
    flex-wrap: wrap;
}

#chatbot .setting-form-max-new-tokens .ms-gr-ant-form-item-control-input-content {
    display: flex;
    flex-wrap: wrap;
}

#chatbot .setting-form-top-p .ms-gr-ant-form-item-control-input-content {
    display: flex;
    flex-wrap: wrap;
}

#chatbot .chatbot-settings-sidebar.hidden {
  display: none;
}
"""

model_options_map_json = json.dumps(MODEL_OPTIONS_MAP)
js = "function init() { window.MODEL_OPTIONS_MAP=" + model_options_map_json + "}"

with gr.Blocks(css=css, js=js, fill_width=True) as demo:
    state = gr.State(
        {
            "conversation_contexts": {},
            "conversations": [],
            "conversation_id": "",
        }
    )

    with (
        ms.Application(),
        antdx.XProvider(theme=DEFAULT_THEME, locale=DEFAULT_LOCALE),
        ms.AutoLoading(),
    ):
        with antd.Row(gutter=[20, 20], wrap=False, elem_id="chatbot"):
            # Left Column
            with antd.Col(
                md=dict(flex="0 0 260px", span=24, order=0),
                span=0,
                elem_style=dict(width=0),
                order=1,
            ):
                with ms.Div(elem_classes="chatbot-conversations"):
                    with antd.Flex(
                        vertical=True,
                        gap="small",
                        elem_style=dict(
                            height="100%",
                            # backgroundImage="linear-gradient(137deg, #e5f4ff 40%, #6400FF 120%)",
                        ),
                    ):
                        # Logo
                        Logo()

                        # New Conversation Button
                        with (
                            antd.Button(
                                value=None,
                                color="primary",
                                variant="solid",  # "outlined", "dashed", "solid", "filled", "text","link"]
                                block=True,
                            ) as add_conversation_btn
                        ):
                            ms.Text("New Chat")
                            with ms.Slot("icon"):
                                antd.Icon("PlusOutlined")

                        # Conversations List
                        with antdx.Conversations(
                            elem_classes="chatbot-conversations-list",
                        ) as conversations:
                            with ms.Slot("menu.items"):
                                with antd.Menu.Item(
                                    label="Delete", key="delete", danger=True
                                ) as conversation_delete_menu_item:
                                    with ms.Slot("icon"):
                                        antd.Icon("DeleteOutlined")
            # Mid Column
            with antd.Col(flex=1, elem_style=dict(height="100%")):
                with antd.Flex(
                    vertical=True,
                    gap="small",
                    elem_classes="chatbot-chat",
                    # elem_style=dict(
                    #     backgroundImage="linear-gradient(137deg, #e5f4ff 10%, #efe7ff 1000%)",
                    # ),
                ):
                    # Chatbot
                    chatbot = pro.Chatbot(
                        elem_classes="chatbot-chat-messages",
                        height=0,
                        welcome_config=welcome_config(),
                        user_config=user_config(),
                        bot_config=bot_config(),
                    )

                    with antdx.Sender(placeholder="Enter your message...") as input:
                        with ms.Slot("prefix"):
                            with antd.Flex(
                                gap=4, wrap=True, elem_style=dict(maxWidth="40vw")
                            ):
                                with antd.Tooltip(title="Clear chat") as tooltip:
                                    with antd.Button(
                                        value=None, type="text"
                                    ) as clear_btn:
                                        with ms.Slot("icon"):
                                            antd.Icon(
                                                "FormatPainterOutlined",
                                                elem_style=dict(color="red"),
                                            )

            # Right Column
            settings_sidebar_visible = gr.State(True)
            with antd.Col(
                md=dict(flex="0 0 355px", span=24, order=2),
                span=0,
                elem_style=dict(width=0),
                order=1,
                visible=True,
            ) as settings_sidebar_col:
                with ms.Div(elem_classes="chatbot-settings-sidebar"):
                    with antd.Flex(
                        vertical=True,
                        gap="small",
                        elem_style=dict(
                            height="100%",
                            padding="32px 16px",
                            # backgroundImage="linear-gradient(137deg, #e5f4ff 10%, #6400FF 115%)",
                        ),
                    ):
                        login_btn = gr.LoginButton()
                        settings_header_state, settings_form = SettingsMenu()

    # Events Handler
    # Browser State Handler
    if save_history:
        browser_state = gr.BrowserState(
            {
                "conversation_contexts": {},
                "conversations": [],
            },
            storage_key="chat_demo_storage",
        )
        state.change(
            fn=Gradio_Events.update_browser_state,
            inputs=[state],
            outputs=[browser_state],
        )

        demo.load(
            fn=Gradio_Events.apply_browser_state,
            inputs=[browser_state, state],
            outputs=[conversations, state],
        )

    # Conversations Handler
    add_conversation_btn.click(
        fn=Gradio_Events.new_chat,
        inputs=[state],
        outputs=[conversations, chatbot, settings_form, state],
    )
    conversations.active_change(
        fn=Gradio_Events.select_conversation,
        inputs=[state],
        outputs=[conversations, chatbot, settings_form, state],
    )
    conversations.menu_click(
        fn=Gradio_Events.click_conversation_menu,
        inputs=[state],
        outputs=[conversations, chatbot, state],
    )
    # Chatbot Handler
    chatbot.welcome_prompt_select(fn=Gradio_Events.apply_prompt, outputs=[input])

    chatbot.delete(fn=Gradio_Events.delete_message, inputs=[state], outputs=[state])
    chatbot.edit(
        fn=Gradio_Events.edit_message, inputs=[state, chatbot], outputs=[state]
    )

    regenerating_event = chatbot.retry(
        fn=Gradio_Events.regenerate_message,
        inputs=[settings_form, state],
        outputs=[
            input,
            clear_btn,
            conversation_delete_menu_item,
            add_conversation_btn,
            conversations,
            chatbot,
            state,
        ],
    )

    def toggle_settings_sidebar(visible):
        return gr.update(value=not visible), gr.update(visible=not visible)

    # Input Handler
    submit_event = input.submit(
        fn=Gradio_Events.add_message,
        inputs=[input, settings_form, state],
        outputs=[
            input,
            clear_btn,
            conversation_delete_menu_item,
            add_conversation_btn,
            conversations,
            chatbot,
            state,
        ],
    )
    input.cancel(
        fn=Gradio_Events.cancel,
        inputs=[state],
        outputs=[
            input,
            conversation_delete_menu_item,
            clear_btn,
            conversations,
            add_conversation_btn,
            chatbot,
            state,
        ],
        cancels=[submit_event, regenerating_event],
        queue=False,
    )

    clear_btn.click(
        fn=Gradio_Events.clear_conversation_history,
        inputs=[state],
        outputs=[chatbot, state],
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=100, max_size=100).launch(
        ssr_mode=False, max_threads=100
    )