import os
import gradio as gr
from huggingface_hub import InferenceClient, login
from custom_tools import suggest_best_cities, RecommendCountryTool
from smolagents import CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, SpeechToTextTool, InferenceClientModel, FinalAnswerTool

login(os.environ.get("HF_TOKEN"))


def create_agent():
    tools = [
        DuckDuckGoSearchTool(),
        WikipediaSearchTool(),
        SpeechToTextTool(),
        suggest_best_cities,
        RecommendCountryTool(),
        FinalAnswerTool()
    ]
    model = InferenceClientModel()
    return CodeAgent(tools=tools, model=model)


def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p
):
    agent = create_agent()

    if isinstance(message, tuple):
        audio_path = message[0]
        stt_tool = SpeechToTextTool()
        message_text = stt_tool.transcribe(audio_path)
        user_message = f"(Transcribed from audio) {message_text}"
    else:
        user_message = message

    full_prompt = f"{system_message}\n\nChat history:\n{history}\n\nUser: {user_message}"

    response = agent.run(full_prompt, stream=False)

    yield response



"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    multimodal=True,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="What's up dude?"),
    ],
)

with gr.Blocks() as demo:
    chatbot.render()


if __name__ == "__main__":
    demo.launch()