#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# 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 mimetypes
import os
import re
import shutil
from typing import Optional
from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
from smolagents.agents import ActionStep, MultiStepAgent
from smolagents.memory import MemoryStep
from smolagents.utils import _is_package_available
import gradio as gr
from gradio.components import Markdown
from gradio.components import Chatbot, Textbox, State, Button

class EnhancedGradioUI:
    """Enhanced Gradio UI with markdown introduction and quick prompt buttons"""

    def __init__(self, agent):
        self.agent = agent

    def interact_with_agent(self, prompt, messages):
        messages.append(gr.ChatMessage(role="user", content=prompt))
        yield messages
        for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
            messages.append(msg)
            yield messages
        yield messages

    def launch(self, **kwargs):
        with gr.Blocks(theme="base") as demo:
            # State to store chat messages
            stored_messages = State([])

            # Markdown Introduction
            gr.Markdown("""
                # NBAi - NBA Stats Chatbot 🤖🏀
                
                Welcome to **NBAi**, your personal NBA statistics assistant! This app fetches and presents NBA box scores from last night's games, giving you insights on player performance, team stats, and more. 

                ## Features
                - Get real-time NBA box scores and player statistics.
                - Ask questions like:
                    - Who had the most points last night?
                    - Who grabbed the most rebounds?
                    - Who had the highest assist-to-turnover ratio?

                ## Tools Used 🔧
                - **smolagents** for building multi-step agents.
                - **Gradio** for the user interface.
                - **BeautifulSoup** and **Pandas** for web scraping and data processing.
                - **DuckDuckGoSearchTool** and **VisitWebpageTool** for enhanced web interactions.

                ## How to Use 🚀
                - Click one of the quick prompt buttons below or type your own question.
                - The chatbot will respond with detailed NBA statistics from last night's games.
                
                ---
            """)

            # Quick Prompt Buttons
            with gr.Row():
                btn_points = Button(value="🏀 Most Points", variant="primary")
                btn_rebounds = Button(value="💪 Most Rebounds", variant="primary")
                btn_assist_to_turnover = Button(value="🎯 Best Assist-to-Turnover Ratio", variant="primary")

            # Chatbot Interface
            chatbot = Chatbot(
                label="NBAi Chatbot",
                type="messages",
                avatar_images=(
                    None,
                    "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",
                ),
                resizeable=True,
                scale=1,
            )

            # Textbox for Custom User Input
            text_input = Textbox(lines=1, label="Your Question")

            # Bindings for Buttons
            btn_points.click(
                self.interact_with_agent,
                ["Who had the most points in last night's NBA games?", stored_messages],
                chatbot
            )

            btn_rebounds.click(
                self.interact_with_agent,
                ["Who had the most rebounds in last night's NBA games?", stored_messages],
                chatbot
            )

            btn_assist_to_turnover.click(
                self.interact_with_agent,
                ["Who had the highest ratio of assists to turnovers in last night's NBA games?", stored_messages],
                chatbot
            )

            # Custom Input Submission
            text_input.submit(
                self.interact_with_agent,
                [text_input, stored_messages],
                chatbot
            )

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