from smolagents import CodeAgent, DuckDuckGoSearchTool, load_tool, tool, LiteLLMModel
import datetime
import requests # Not used in the provided snippet directly, but often a common import
import pytz # Used in the commented-out example tool
import yaml
import os
import tempfile
from tools.final_answer import FinalAnswerTool # Assuming this exists in ./tools/final_answer.py
from tools.visit_webpage import VisitWebpageTool # Assuming this exists in ./tools/visit_webpage.py
from tools.web_search import DuckDuckGoSearchTool
from smolagents import GradioUI
import gradio as gr
import json # Not used directly in the provided snippet, but can be common
import os # os is imported twice, which is fine.
from typing import Dict, List, Optional, Union, Any
import re # Added for regex in Gradio UI if not already globally available

'''
# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_custom_tool(x:str, y:int)-> int: #it's import to specify the return type
    #Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"
@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"
'''
@tool
def create_document(text: str, format: str = "docx") -> str:
    """Creates a document with the provided text and allows download.
    
    Args:
        text: The text content to write to the document.
        format: The output format, either 'docx', 'pdf', or 'txt'.
    """
    try:
        # Create a temp directory to store files
        temp_dir = tempfile.mkdtemp()
        
        if format.lower() == "txt":
            txt_path = os.path.join(temp_dir, "generated_document.txt")
            with open(txt_path, "w", encoding="utf-8") as f:
                f.write(text)
            return txt_path

        elif format.lower() in ["docx", "pdf"]:
            try:
                import docx
                from docx.shared import Pt
            except ImportError:
                return (f"ERROR: To create DOCX or PDF files, the 'python-docx' package is required. "
                        f"Please install it (e.g., 'pip install python-docx'). "
                        f"You can try creating a 'txt' file instead by specifying format='txt'.")

            # Create a new document
            doc = docx.Document()
            doc.add_heading('Generated Document', 0)
            style = doc.styles['Normal']
            font = style.font
            font.name = 'Calibri'
            font.size = Pt(11)
            
            for paragraph in text.split('\n'):
                if paragraph.strip():  # Skip empty paragraphs
                    doc.add_paragraph(paragraph)
            
            docx_path = os.path.join(temp_dir, "generated_document.docx")
            doc.save(docx_path)
            
            if format.lower() == "pdf":
                try:
                    from docx2pdf import convert
                    pdf_path = os.path.join(temp_dir, "generated_document.pdf")
                    convert(docx_path, pdf_path)
                    # Optional: os.remove(docx_path) # If PDF is main target and DOCX is intermediate
                    return pdf_path
                except ImportError:
                    return (f"ERROR: PDF conversion requires the 'docx2pdf' package. "
                            f"Please install it (e.g., 'pip install docx2pdf'). "
                            f"Document saved as DOCX instead at: {docx_path}")
                except Exception as e_pdf:
                    return f"Error converting DOCX to PDF: {str(e_pdf)}. Document saved as DOCX at: {docx_path}"
            
            return docx_path
        
        else:
            return f"Error: Unsupported format '{format}'. Supported formats are 'docx', 'pdf', 'txt'."
            
    except Exception as e:
        # Catch-all for other unexpected errors during file operations, etc.
        return f"Error creating document: {str(e)}"

# Custom file download tool to help with file handling
@tool
def get_file_download_link(file_path: str) -> str:
    """Creates a download link for a file.
    
    Args:
        file_path: Path to the file that should be made available for download
    """
    if not os.path.exists(file_path):
        return f"Error: File not found at {file_path}"
    
    # Get file extension and set up appropriate mime type
    _, file_extension = os.path.splitext(file_path)
    mime_types = {
        '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
        '.pdf': 'application/pdf',
        '.txt': 'text/plain',  # Added for .txt files
    }
    mime_type = mime_types.get(file_extension.lower(), 'application/octet-stream')
    
    # Return information that can be used by the agent to instruct the user
    return f"File ready for download: {os.path.basename(file_path)} ({mime_type})"
        

final_answer = FinalAnswerTool()
web_search=DuckDuckGoSearchTool() # This was commented out in original, keeping it so
visit_webpage=VisitWebpageTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = LiteLLMModel(
    model_id="gemini/gemini-2.0-flash", # As per original, though gemini-1.5-flash-latest might be more common
    api_key=os.getenv("GEMINI_KEY"),
    max_tokens=8192,
    temperature=0.7
)


# Import tool from Hub
# image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True) # This tool is loaded but not added to the agent

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

#,web_search
    
agent = CodeAgent(
    model=model,
    tools=[final_answer,visit_webpage,create_document,get_file_download_link,web_search], ## add your tools here (don't remove final answer)
    max_steps=10,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

# Custom Gradio UI with file download capability
class CustomGradioUI(GradioUI):
    def build_interface(self):
        with gr.Blocks() as interface:
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("# AI Assistant")
                
            chatbot = gr.Chatbot(height=600)
            msg = gr.Textbox(
                placeholder="Ask me anything...",
                container=False,
                scale=7,
            )
            
            # Add a file download component
            download_btn = gr.Button("Download File", visible=False)
            file_output = gr.File(label="Generated Document", visible=False)
            
            # Store the latest file path
            self._latest_file_path = None # Instance variable to store path between calls
            
            def respond(message, chat_history):
                agent_response = self.agent.run(message)
                chat_history.append((message, agent_response))
                
                show_download = False
                # Reset latest_file_path to ensure it's fresh for each response
                # self._latest_file_path = None # Actually, should be set if file is found

                # Look for generated file paths in the agent's response.
                # tempfile.mkdtemp() usually creates paths like /tmp/some_random_string/
                # The tool consistently names files "generated_document.ext"
                paths = re.findall(r'/tmp/[^/]+/generated_document\.(docx|pdf|txt)', agent_response)
                if paths:
                    self._latest_file_path = paths[0] # Store the full path
                    show_download = True
                else:
                    self._latest_file_path = None # Clear if no path found in this response
                
                return chat_history, gr.Button.update(visible=show_download), gr.File.update(visible=False) # Keep file output hidden until download clicked
            
            def prepare_download():
                if self._latest_file_path and os.path.exists(self._latest_file_path):
                    return gr.File.update(value=self._latest_file_path, visible=True)
                # Handle case where file might have been cleaned up or path is stale
                gr.Warning("File not available for download. It might have been temporary or an error occurred.")
                return gr.File.update(visible=False)
            
            # Connect the components
            msg.submit(respond, [msg, chatbot], [chatbot, download_btn, file_output])
            download_btn.click(prepare_download, [], [file_output])
            
            gr.Markdown("Powered by smolagents and Qwen") # Qwen is mentioned, perhaps the HF endpoint uses it.
            
        return interface

# Ensure this conditional execution if the script can also be imported
if __name__ == "__main__":
    # Note: The original script had GradioUI(agent).launch() directly.
    # Using the custom UI:
    CustomGradioUI(agent).launch()