from smolagents import CodeAgent, HfApiModel, tool as smol_tool, GradioUI
from smolagents.tools import Tool
import datetime
import pytz
import yaml
import tempfile
import gradio as gr
import os
import re

# Import the tools from your local folder
from tools.final_answer import FinalAnswerTool
from tools.visit_webpage import VisitWebpageTool
from tools.create_presentation import create_presentation
from tools.web_search import DuckDuckGoSearchTool
from tools.write_code import write_code_to_file
# =========================
# ADDITIONAL IN-LINE TOOLS
# =========================

@smol_tool
def create_document(text: str, format: str = "docx") -> str:
    """
    Creates a document file (DOCX or PDF) containing the provided text.
    Args:
        text: The full textual content that will be written into the document.
        format: The output file format. Must be either "docx" or "pdf".
    """
    try:
        import docx
        from docx.shared import Pt
        temp_dir = tempfile.mkdtemp()
        doc = docx.Document()
        doc.add_heading("Generated Document", 0)
        
        for paragraph in text.split("\n"):
            if paragraph.strip():
                doc.add_paragraph(paragraph)

        file_path = os.path.join(temp_dir, f"generated_document.{format.lower()}")
        doc.save(file_path)
        return file_path
    except Exception as e:
        return f"Error creating document: {str(e)}"

@smol_tool
def create_presentation_preview(title: str, slides: list[str]) -> list[str]:
    """
    Create a presentation and return slide previews as images.
    Args:
        title: Presentation title
        slides: List of slide contents
    """
    from PIL import Image, ImageDraw
    temp_dir = tempfile.mkdtemp()
    image_paths = []
    for i, slide_text in enumerate(slides, start=1):
        img = Image.new("RGB", (800, 450), "white")
        draw = ImageDraw.Draw(img)
        draw.text((40, 40), f"{title}\n\nSlide {i}:\n{slide_text[:100]}", fill="black")
        path = os.path.join(temp_dir, f"slide_{i}.png")
        img.save(path)
        image_paths.append(path)
    return image_paths

# =========================
# AGENT SETUP
# =========================
final_answer = FinalAnswerTool()
visit_webpage = VisitWebpageTool()
web_search = DuckDuckGoSearchTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id="Qwen/Qwen2.5-72B-Instruct", # Updated to a widely available stable ID
)

# Load the prompt templates we created earlier
with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[
        final_answer,
        visit_webpage,
        web_search, # Added from your tools folder
        create_document,
        create_presentation,
        create_presentation_preview,
        write_code_to_file,
    ],
    max_steps=10,
    verbosity_level=1,
    prompt_templates=prompt_templates,
)

# =========================
# UI CUSTOMIZATION
# =========================
class CustomGradioUI(GradioUI):
    def build_interface(self):
        with gr.Blocks(theme=gr.themes.Soft()) as interface:
            gr.Markdown("# 🚀 Enterprise AI Agent")
            gr.Markdown("Search the web, create docs, and generate PowerPoints instantly.")

            chatbot = gr.Chatbot(height=500, show_label=False)
            with gr.Row():
                msg = gr.Textbox(placeholder="E.g., 'Research AI trends and make a 5-slide PPT'", scale=9)
                submit = gr.Button("Send", variant="primary", scale=1)

            with gr.Row():
                download_btn = gr.Button("📂 Download Generated File", visible=False)
                file_output = gr.File(label="Ready for Download", visible=False)

            self._latest_file_path = None

            def respond(message, chat_history):
                # Run the agent
                agent_response = self.agent.run(message)
                chat_history.append((message, str(agent_response)))

                # File detection logic
                self._latest_file_path = None
                show_download = False
                match = re.search(r"(/tmp/[\w\d_-]+/generated_[\w-]+\.(docx|pdf|pptx))", str(agent_response))
                
                if match:
                    self._latest_file_path = match.group(1)
                    show_download = True

                return chat_history, gr.update(visible=show_download), gr.update(visible=False)

            def prepare_download():
                if self._latest_file_path and os.path.exists(self._latest_file_path):
                    return gr.update(value=self._latest_file_path, visible=True)
                return gr.update(visible=False)

            msg.submit(respond, [msg, chatbot], [chatbot, download_btn, file_output])
            submit.click(respond, [msg, chatbot], [chatbot, download_btn, file_output])
            download_btn.click(prepare_download, [], [file_output])

        return interface

if __name__ == "__main__":
    CustomGradioUI(agent).launch()