# streamlit helpers for ChatAgent
import datetime
import json
import time
from base64 import b64decode, b64encode
from typing import TYPE_CHECKING, Callable, Optional

import streamlit as st
from humanize import precisedelta
from pydantic import BaseModel, ValidationError
from streamlit.runtime.uploaded_file_manager import UploadedFile

from supersullytools.llm.agent import AgentTool, ChatAgent
from supersullytools.llm.completions import CompletionHandler, CompletionModel, ImagePromptMessage, PromptAndResponse
from supersullytools.llm.trackers import SessionUsageTracking, StoredPromptAndResponse
from supersullytools.streamlit.paginator import item_paginator
from supersullytools.utils.misc import format_validation_error

if TYPE_CHECKING:
    from supersullytools.utils.media_manager import MediaManager


class SlashCmd(BaseModel):
    name: str
    description: str
    mechanism: Callable
    refresh_after: bool = False


class ChatAgentUtils(object):
    def __init__(
        self,
        chat_agent: ChatAgent,
        use_system_slash_cmds: bool = True,
        extra_slash_cmds: Optional[dict[str, SlashCmd]] = None,
    ):
        self.chat_agent = chat_agent
        self.use_system_slash_cmds = use_system_slash_cmds
        self._system_slash_cmds = {
            "/help": SlashCmd(
                name="Help",
                description="Display available commands",
                mechanism=self._display_slash_help,
                refresh_after=False,
            ),
        }

        if self.use_system_slash_cmds:
            self._system_slash_cmds["/clear"] = SlashCmd(
                name="Clear Chat",
                description="Clear current chat history",
                mechanism=self.chat_agent.reset_history,
                refresh_after=True,
            )
            self._system_slash_cmds["/tool"] = SlashCmd(
                name="Use a Tool",
                description="Opens a dialog to enable manual tool usage",
                mechanism=self._use_tool_manually,
                refresh_after=False,
            )
            if self.has_session_tracker():
                self._system_slash_cmds["/completions"] = SlashCmd(
                    name="View LLM Completions",
                    description="Opens a dialog to browse LLM completions from the session",
                    mechanism=self._browse_session_completions,
                    refresh_after=False,
                )
        self.extra_slash_cmds = extra_slash_cmds or {}

    def get_session_tracker(self) -> Optional[SessionUsageTracking]:
        if not self.chat_agent.completion_handler.completion_tracker:
            return None
        trackers = self.chat_agent.completion_handler.completion_tracker.trackers
        try:
            return next(x for x in trackers if isinstance(x, SessionUsageTracking))
        except StopIteration:
            return None

    def has_session_tracker(self) -> bool:
        return bool(self.get_session_tracker())

    def _display_slash_help(self):
        output = "### Available Commands\n\n"
        all_commands = {**self.extra_slash_cmds, **self._system_slash_cmds}
        for slash_cmd, obj in all_commands.items():
            obj: SlashCmd
            output += f"* **{obj.name}** (`{slash_cmd}`): {obj.description}\n"
        return output

    def get_completion_model(self, model: Optional[str | CompletionModel] = None) -> CompletionModel:
        if isinstance(model, CompletionModel):
            return model
        return self.chat_agent.completion_handler.get_model_by_name_or_id(model)

    @staticmethod
    def select_llm(
        completion_handler: CompletionHandler, label, default: str = "GPT 4 Omni Mini", key=None, **kwargs
    ) -> CompletionModel:
        completion_handler = completion_handler
        default_model = completion_handler.get_model_by_name_or_id(default)
        return st.selectbox(
            label,
            completion_handler.available_models,
            completion_handler.available_models.index(default_model),
            format_func=lambda x: x.llm,
            key=key,
            **kwargs,
        )

    def select_llm_from_agent(
        self, label, default_override: Optional[str] = None, key=None, **kwargs
    ) -> CompletionModel:
        default_model = (
            self.chat_agent.completion_handler.get_model_by_name_or_id(default_override)
            if default_override
            else self.chat_agent.default_completion_model
        )
        return self.select_llm(
            completion_handler=self.chat_agent.completion_handler,
            label=label,
            default=default_model.llm,
            key=key,
            **kwargs,
        )

    def display_chat_msg(self, msg: str):
        if "<tool>" in msg:
            content, _ = msg.split("<tool>", maxsplit=1)
            content = content.strip()
            try:
                tool_calls = self.chat_agent.extract_tool_calls_from_msg(msg)
            except Exception:
                content += "\n\n<msg contains malformed tool call>"
                tool_calls = []
            if content:
                st.write(content)
            for tc in tool_calls:
                if params := tc.get("parameters"):
                    with st.popover(tc["name"]):
                        st.write(params)
                else:
                    st.caption(f"used `{tc['name']}`")
        elif "<tool_result>" in msg:
            _, result_str = msg.split("<tool_result>", maxsplit=1)
            tool_result_str = [
                this_result_str.split("</tool_result>", maxsplit=1)[0]
                for this_result_str in result_str.split("<tool_result>")
            ]
            tool_results = [x.strip() for x in tool_result_str]
            st.json(tool_results, expanded=0)

        else:
            st.write(msg)

    def _try_handle_system_slash_command(self, command_str: str) -> Optional[tuple[SlashCmd, str]]:
        if command_str in self._system_slash_cmds:
            command: SlashCmd = self._system_slash_cmds[command_str]
            output = command.mechanism()
            if output:
                with st.chat_message("system").container(border=True):
                    st.write(output)
            return command, output

    def _try_handle_extra_slash_cmd(self, command_str: str) -> Optional[tuple[SlashCmd, str]]:
        if command_str in self.extra_slash_cmds:
            command: SlashCmd = self.extra_slash_cmds[command_str]
            output = command.mechanism()
            if output:
                with st.chat_message("system").container(border=True):
                    st.write(output)
            return command, output

    def add_user_message(self, msg: str, images: Optional[list[UploadedFile]] = None) -> bool:
        """Returns true if the streamlit app should reload."""
        if msg.startswith("/"):
            if not (self.use_system_slash_cmds or self.extra_slash_cmds):
                raise RuntimeError("No slash commands available!")
            executed_command = None
            if self.use_system_slash_cmds or msg.startswith("/help"):
                result = self._try_handle_system_slash_command(msg)
                if result:
                    executed_command, _ = result
            if not executed_command and self.extra_slash_cmds:
                result = self._try_handle_extra_slash_cmd(msg)
                if result:
                    executed_command, _ = result
            if not executed_command:
                executed_command, _ = self._try_handle_system_slash_command("/help")
            return executed_command.refresh_after
        else:
            if images:
                prompt = ImagePromptMessage(
                    content=msg,
                    images=[b64encode(image.getvalue()).decode() for image in images],
                    image_formats=["png" if image.name.endswith("png") else "jpeg" for image in images],  # noqa
                )
            else:
                prompt = msg
            self.chat_agent.message_from_user(prompt)
            return True

    def display_chat_and_run_agent(self, include_function_calls):
        num_chat_before = len(self.chat_agent.get_chat_history(include_function_calls=include_function_calls))

        for msg in self.chat_agent.get_chat_history(include_function_calls=include_function_calls):
            with st.chat_message(msg.role):
                if isinstance(msg, ImagePromptMessage):
                    main, images = st.columns((90, 10))
                    with main:
                        self.display_chat_msg(msg.content)
                    for x in msg.images:
                        images.image(b64decode(x.encode()))
                else:
                    self.display_chat_msg(msg.content)

        if self.chat_agent.working:
            with st.status("Agent working...", expanded=True) as status:
                # Define the callback function within the scope of `status`
                def status_callback_fn(message):
                    status.update(label=f"Agent working... {message}", state="running")
                    st.write(message)

                # Run the agent loop, passing the callback function
                while self.chat_agent.working:
                    self.chat_agent.run_agent(status_callback_fn=status_callback_fn)
                    time.sleep(0.05)

                # Final status update when the agent completes
                status.update(label="Agent completed work!", state="complete", expanded=False)

        # output any new messages
        for msg in self.chat_agent.get_chat_history(include_function_calls=include_function_calls)[num_chat_before:]:
            with st.chat_message(msg.role):
                self.display_chat_msg(msg.content)

    def _browse_session_completions(self):
        now = datetime.datetime.now(tz=datetime.timezone.utc)

        @st.dialog("Completions Viewer", width="large")
        def _d():
            if not self.has_session_tracker():
                st.warning("No session tracker found!")

            st.caption("Completions displayed newest-first")
            completions = list(reversed(self.get_session_tracker().completions))

            def _display(idx):
                display_completion(completions[idx], now)

            item_paginator(
                "Completion",
                [x.prompt[-1].content[:25] for x in completions],
                item_handler_fn=_display,
                enable_keypress_nav=True,
                display_item_names=True,
            )

        return _d()

    def _use_tool_manually(self):
        @st.dialog("Trigger Tool", width="large")
        def _d():
            tool_to_use: AgentTool = st.selectbox(
                "Tool to use", self.chat_agent.get_current_tools(), format_func=lambda x: x.name
            )
            tool_dict = self.chat_agent.tool_description_to_dict(tool_to_use)
            if descr := tool_dict.get("description"):
                st.write(descr)
            st.write(tool_dict["parameters"])
            default = {x: None for x in tool_dict["parameters"]}
            params_str = st.text_area("tool params", json.dumps(default, indent=2), height=300)

            try:
                params_obj = tool_to_use.params_model.model_validate_json(params_str)
            except ValidationError as e:
                st.error(format_validation_error(e))
                with st.popover("Invoke Tool", disabled=True):
                    pass
            else:
                invoke_params = json.loads(params_str)
                with st.popover("Invoke Tool"):
                    st.subheader(f"Calling tool {tool_to_use.name}")
                    st.write("Tool Call Parmeters:")
                    st.code(params_obj.model_dump_json(indent=2))
                    if st.button(
                        "Invoke Locally",
                        help="Execute the tool and show the results here, but do to store in session/chat",
                    ):
                        st.write(tool_to_use.invoke_tool(invoke_params))
                    if st.button("Have Agent Invoke"):
                        self.chat_agent.manually_invoke_tool(tool_to_use.name, invoke_params)
                        time.sleep(0.01)
                        st.rerun()

        return _d()


@st.cache_data
def get_media_preview(_media_manager: "MediaManager", media_id):
    return _media_manager.retrieve_media_preview(media_id)


@st.cache_data
def get_stored_media(_media_manager: "MediaManager", media_id):
    return _media_manager.retrieve_media_contents(media_id)


def display_completion(
    par: PromptAndResponse | StoredPromptAndResponse, now, media_manager: Optional["MediaManager"] = None
):
    generated_ago = precisedelta(now - par.response.generated_at, minimum_unit="minutes")
    st.caption(par.response.generated_at.isoformat() + f" ({generated_ago} ago)")

    get_full_image_contents = False
    if media_manager and isinstance(par, StoredPromptAndResponse):
        get_full_image_contents = st.toggle("Retrieve full Image contents")

    is_stored_par = isinstance(par, StoredPromptAndResponse)

    if is_stored_par:
        st.metric("Stored Size", par.get_db_item_size())
        par.resource_config["compress_data"] = False
        st.metric("Uncompressed Size", par.get_db_item_size())
        par.resource_config["compress_data"] = True

    if st.toggle("Show raw"):
        if is_stored_par:
            st.code(par.model_dump_json(indent=2, exclude=par.get_db_resource_base_keys()))
        else:
            st.code(par.model_dump_json(indent=2))
    else:
        st.write("**Prompt**")
        with st.container(border=True):
            if isinstance(par.prompt, str):
                st.code(par.prompt)
            else:
                for idx, msg in enumerate(par.prompt):
                    with st.chat_message(msg.role):
                        if len(msg.content) > 100:
                            with st.popover(msg.content[:100] + " ... ", use_container_width=True):
                                st.write(msg.content)
                        else:
                            st.write(msg.content)
                        if isinstance(msg, ImagePromptMessage):
                            cols = iter(st.columns(len(msg.images)))
                            for x in msg.images:
                                next(cols).image(b64decode(x.encode()))
                        if is_stored_par:
                            if idx in par.prompt_image_media_ids:
                                st.caption("This message included image content")
                                stored_media_ids = par.prompt_image_media_ids[idx]
                                if not stored_media_ids:
                                    st.warning("Images not preserved")
                                    continue
                                if media_manager:
                                    cols = iter(st.columns(len(stored_media_ids)))
                                    for x in stored_media_ids:
                                        if get_full_image_contents:
                                            next(cols).image(get_stored_media(media_manager, x))
                                        else:
                                            with next(cols):
                                                st.image(get_media_preview(media_manager, x))
                                                load_full = st.button("➕", key=x)
                                            if load_full:
                                                st.image(get_stored_media(media_manager, x), use_column_width=True)

        st.write("**Response**")
        with st.container(border=True):
            st.code(par.response.content)

        st.write("**Metadata**")
        with st.container(border=True):
            cols = iter(st.columns(2))
            with next(cols):
                st.json(par.response.llm_metadata.model_dump(mode="json"))
            with next(cols):
                st.json(par.response.model_dump(mode="json", exclude={"content", "llm_metadata", "response_metadata"}))
            st.write("**Raw Provider Response Metadata**")
            st.json(par.response.response_metadata or {})

    st.caption(par.response.generated_at.isoformat() + f" ({generated_ago} ago)")
