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

import gradio as gr
import json # Not used directly in the provided snippet, but can be common

from typing import Dict, List, Optional, Union, Any
import re # Added for regex in Gradio UI if not already globally available

from smolagents.tools import Tool as SmolToolBase # Renamed to avoid conflict with @tool decorator
from smolagents.agent_types import AgentAudio, AgentImage, AgentText # For GradioUI output handling

from Gradio_UI import GradioUI # This will import the class from your modified file


@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'. Default is 'docx'.
    """
    try:
        temp_dir = tempfile.mkdtemp()
        file_name = "generated_document" # Consistent filename
        
        if format.lower() == "txt":
            path = os.path.join(temp_dir, f"{file_name}.txt")
            with open(path, "w", encoding="utf-8") as f:
                f.write(text)
            print(f"Document created (txt): {path}")
            return 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'.")

            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():
                    doc.add_paragraph(paragraph)
            
            docx_path = os.path.join(temp_dir, f"{file_name}.docx")
            doc.save(docx_path)
            print(f"Document created (docx): {docx_path}")

            if format.lower() == "pdf":
                try:
                    from docx2pdf import convert
                    pdf_path = os.path.join(temp_dir, f"{file_name}.pdf")
                    convert(docx_path, pdf_path)
                    print(f"Document converted to PDF: {pdf_path}")
                    return pdf_path
                except ImportError:
                    err_msg = (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}")
                    print(err_msg)
                    return err_msg # Or return docx_path with a note
                except Exception as e_pdf:
                    err_msg = f"Error converting DOCX to PDF: {str(e_pdf)}. Document saved as DOCX at: {docx_path}"
                    print(err_msg)
                    return err_msg # Or return docx_path
            
            return docx_path
        
        else:
            return f"Error: Unsupported format '{format}'. Supported formats are 'docx', 'pdf', 'txt'."
            
    except Exception as e:
        print(f"General error in create_document: {str(e)}")
        return f"Error creating document: {str(e)}"

@tool
def get_file_download_link(file_path: str) -> str:
    """Informs that a file is ready for download and its type. (Used by agent to tell user).
    Args:
        file_path: Path to the file that should be made available for download.
    """
    if not os.path.exists(file_path):
        print(f"get_file_download_link: File not found at {file_path}")
        return f"Error: File not found at {file_path}"
    
    _, file_extension = os.path.splitext(file_path)
    mime_types = {
        '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
        '.pdf': 'application/pdf',
        '.txt': 'text/plain',
    }
    mime_type = mime_types.get(file_extension.lower(), 'application/octet-stream')
    msg = f"File '{os.path.basename(file_path)}' is ready for download (type: {mime_type})."
    print(f"get_file_download_link: {msg}")
    return msg
        
final_answer = FinalAnswerTool()
web_search=DuckDuckGoSearchTool()
visit_webpage=VisitWebpageTool()

# Load image generation tool from Hub
try:
    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True) # More likely correct
    # Common alternative: "agents-course/text-to-image" (uses a smaller SD model)
    # image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
    print(f"Successfully loaded image generation tool. Name: {getattr(image_generation_tool, 'name', 'N/A')}")
except Exception as e:
    print(f"Error loading image generation tool: {e}")
    print("Image generation will not be available.")
    image_generation_tool = None # Set to None if loading fails

# --- Configure LLM Model ---
# Ensure GEMINI_KEY is set in your environment variables
if not os.getenv("GEMINI_KEY"):
    print("CRITICAL: GEMINI_KEY environment variable not set. The LLM will not work.")
    # exit(1) # Or handle gracefully

model = LiteLLMModel(
    model_id="gemini/gemini-1.5-flash-latest", # Using a common and capable Gemini model
    api_key=os.getenv("GEMINI_KEY"),
    max_tokens=4096, # Max output tokens for the LLM
    temperature=0.6 # Slightly lower temperature for more factual tool use
)

# --- Load Prompts ---
prompts_file = "prompts.yaml"
if not os.path.exists(prompts_file):
    print(f"Warning: '{prompts_file}' not found. Using default agent prompts.")
    prompt_templates = None # Agent will use its defaults
else:
    with open(prompts_file, 'r') as stream:
        prompt_templates = yaml.safe_load(stream)

# --- Initialize Agent ---
agent_tools = [
    final_answer,
    web_search,
    visit_webpage,
    create_document, # This is a @tool decorated function, so it's directly usable
    get_file_download_link, # Also a @tool decorated function
]

if image_generation_tool:
    agent_tools.append(image_generation_tool)
else:
    print("Note: Image generation tool was not loaded, so it's not added to the agent.")
    
agent = CodeAgent(
    model=model,
    tools=agent_tools,
    max_steps=10, # Max iterations for the agent
    verbosity_level=1, # Set to 1 or higher to populate interaction_logs for GradioUI
    prompt_templates=prompt_templates
)

# --- Main Application Entry Point ---
if __name__ == "__main__":
    print("Starting Gradio UI...")
    # Initialize GradioUI with the agent
    # The GradioUI class should be imported from your modified Gradio_UI.py
    # Ensure Gradio_UI.py is in the same directory or python path
    ui = GradioUI(agent, file_upload_folder="uploaded_files") # Enable file uploads
    ui.launch() # share=False is default and safer for Spaces