# Copyright 2024-2025 Akihito Miyazaki.
# This code is derived from the DuckDuckGoSearchTool class,
# originally part of the HuggingFace smolagents library.
# https://github.com/huggingface/smolagents
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import os
import threading
from typing import Optional

from dotenv import load_dotenv
import gradio as gr
from smolagents import (
    CodeAgent,
    LiteLLMModel,
    DuckDuckGoSearchTool,
)
from smolagents.agent_types import AgentText, AgentImage, AgentAudio
from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
from huggingface_hub import InferenceClient


from extra_search_tools import (
    PrioritySearchTool,
    BraveSearchTool,
    GoogleCustomSearchTool,
)


def hf_chat(api_key, model, text):
    client = InferenceClient(api_key=api_key)
    messages = [
        {
            "role": "user",
            "content": text,
        }
    ]

    stream = client.chat.completions.create(
        model=model, messages=messages, max_tokens=6000, stream=False
    )

    return stream.choices[0].message.content


load_dotenv(override=True)
# login(os.getenv("HF_TOKEN"))

append_answer_lock = threading.Lock()

# without below chat will duplicate
custom_role_conversions = {"tool-call": "assistant", "tool-response": "user"}

model = LiteLLMModel(
    "groq/llama3-8b-8192",
    api_base="https://api.groq.com/openai/v1",
    max_completion_tokens=500,
    api_key=os.getenv("GROQ_API_KEY"),  # Groq API
)

search_tool = PrioritySearchTool(
    [
        DuckDuckGoSearchTool(),
        GoogleCustomSearchTool("YOUR_ENGINE_KEY"),
        BraveSearchTool(),
    ],
    "history.json",
)
WEB_TOOLS = [search_tool]

max_steps = 1


# Agent creation in a factory function
def create_agent():
    print("create agent")
    """Creates a fresh agent instance for each session"""
    return CodeAgent(
        model=model,
        tools=WEB_TOOLS,
        max_steps=max_steps,
        verbosity_level=1,
    )


def stream_to_gradio(
    agent,
    task: str,
    reset_agent_memory: bool = False,
    additional_args: Optional[dict] = None,
):
    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
    steps = 0
    for step_log in agent.run(
        task, stream=True, reset=reset_agent_memory, additional_args=additional_args
    ):
        # I dont know the reason but call more steps
        steps += 1
        if steps <= max_steps:
            for message in pull_messages_from_step(
                step_log,
            ):
                yield message

    final_answer = step_log  # Last log is the run's final_answer
    final_answer = handle_agent_output_types(final_answer)
    # print(final_answer)
    if isinstance(final_answer, AgentText):
        yield gr.ChatMessage(
            role="assistant",
            content=f"**Final answer:**\n{final_answer.to_string()}",
        )
    elif isinstance(final_answer, AgentImage):
        yield gr.ChatMessage(
            role="assistant",
            content={"path": final_answer.to_string(), "mime_type": "image/png"},
        )
    elif isinstance(final_answer, AgentAudio):
        yield gr.ChatMessage(
            role="assistant",
            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
        )
    else:
        yield gr.ChatMessage(
            role="assistant", content=f"**Final answer:** {str(final_answer)}"
        )


class GradioUI:
    """A one-line interface to launch your agent in Gradio"""

    def __init__(self, file_upload_folder: str | None = None):
        self.file_upload_folder = file_upload_folder
        if self.file_upload_folder is not None:
            if not os.path.exists(file_upload_folder):
                os.mkdir(file_upload_folder)

    def interact_with_agent(self, prompt, messages, session_state):
        # Get or create session-specific agent
        if "agent" not in session_state:
            session_state["agent"] = create_agent()

        messages.append(gr.ChatMessage(role="user", content=prompt))
        yield messages

        # Use session's agent instance
        for msg in stream_to_gradio(
            session_state["agent"], task=prompt, reset_agent_memory=False
        ):
            messages.append(msg)
            pass
            yield messages
        yield messages

    def launch(self, **kwargs):
        with gr.Blocks(theme="ocean", fill_height=True) as demo:
            gr.Markdown("""# Smolagents - ExtraSearchtools!
- [Google Custom Search](https://developers.google.com/custom-search/v1/overview) tool
- [Brave Search](https://brave.com/search/api/) tool
- PrioritySearchTool - try duckduckgo fist and then use google
- PrioritySearchTool - json-save function
                        
Built with [smolagents](https://github.com/huggingface/smolagents).This Demo only work duckduckgo if it's not rate-limited.Duplicate and set your own secret key 
""")
            # Add session state to store session-specific data
            session_state = gr.State({})  # Initialize empty state for each session

            chatbot = gr.Chatbot(
                label="ExtraSearchtools",
                type="messages",
                avatar_images=(
                    None,
                    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
                ),
                scale=1,
            )
            text_input = gr.Textbox(
                lines=1, label="Your request", value="What is smolagents?"
            )
            text_input.submit(
                self.interact_with_agent,
                # Include session_state in function calls
                [text_input, chatbot, gr.State({})],
                [chatbot],
            )

        demo.launch(debug=True, share=True, **kwargs)


if __name__ == "__main__":
    GradioUI().launch()  # not support auto update restart by yourself :AttributeError: module '__main__' has no attribute 'demo'