import os
from docx import Document
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel,load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_custom_tool(arg1:str, arg2:int)-> str: #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
from docx import Document

class UiPathReportTool:
    def __init__(self):
        pass

    def generate_automation_report(self, case_description: str) -> str:
        """A tool that generates a structured Word report for RPA analysts based on a case description."""
        # Create a new Word document
        doc = Document()
        
        # Report title
        doc.add_heading('Automation Suitability & Requirements Report', 0)
        
        # Case Description
        doc.add_heading('Case Description', 1)
        p = doc.add_paragraph(case_description)
        p.style = 'List Bullet'  # Example formatting
        
        # Process Assessment Tool
        doc.add_heading('Process Assessment Tool (PAT) for Automation', 1)
        pat_table = doc.add_table(rows=1, cols=3)
        hdr_cells = pat_table.rows[0].cells
        hdr_cells[0].text = 'METRIC'
        hdr_cells[1].text = 'RATING (1-5)'
        hdr_cells[2].text = 'JUSTIFICATION'
        # Example assessment
        row = pat_table.add_row().cells
        row[0].text = 'Frequency of Execution'
        row[1].text = '5'
        row[2].text = 'Process is executed daily with predictable patterns'
        
        # Requirements Template
        doc.add_heading('Business Requirements Template', 1)
        doc.add_heading('I. PROJECT SCOPE', 2)
        doc.add_paragraph('The goal of this automation is to... ')
        # Add requirement sections here (Objectives, Constraints, Compliance, etc.)
        
        # PDD (Process Definition Document)
        doc.add_heading('Process Definition Document (PDD)', 1)
        doc.add_heading('A. Process Overview', 2)
        doc.add_paragraph('This section describes the current manual process.')
        # Add PDD sections like Input/Output Data, Process Step Diagram, Stakeholders
        
        # SDD (System Design Document)
        doc.add_heading('System Design Document (SDD)', 1)
        doc.add_heading('1. Technical Architecture', 2)
        doc.add_paragraph('This section outlines the system architecture.')
        # Add SDD sections (Integration Points, Error Handling, Validation Rules)
        
        # UAT (User Acceptance Testing)
        doc.add_heading('User Acceptance Testing (UAT) Plan', 1)
        doc.add_heading('A. Test Cases', 2)
        uat_table = doc.add_table(rows=1, cols=4)
        hdr_cells = uat_table.rows[0].cells
        hdr_cells[0].text = 'TEST ID'
        hdr_cells[1].text = 'CASE DESCRIPTION'
        hdr_cells[2].text = 'EXPECTED RESULT'
        hdr_cells[3].text = 'STATUS'
        # Example test case
        row = uat_table.add_row().cells
        row[0].text = 'UAT-001'
        row[1].text = 'Validate data extraction accuracy'
        row[2].text = 'All fields correctly populated'
        row[3].text = 'PASSED'  # Default placeholder
        
        # Analysis Charts
        doc.add_heading('Analysis Charts & Key Metrics', 1)
        doc.add_paragraph('Include charts like Process Complexity vs. Frequency, ROI Projection.')
        doc.add_paragraph('These visualizations can be inserted as Excel graphs or Visio diagrams.')
        
        # Save the document
        report_filename = 'Automation_Suitability_Report.docx'
        doc.save(report_filename)
        
        return (f"Automation suitability report successfully generated: {report_filename}. "
                f"Replace placeholders with project-specific details and add supporting images as needed.")

# Example usage
if __name__ == "__main__":
    report_tool = UiPathReportTool()
    case_desc = "A manual process where sales orders are entered into ERP from emails, currently taking 30 minutes per order."
    result = report_tool.generate_automation_report(case_desc)
    print(result)

# 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 = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


    
agent = CodeAgent(
    model=model,
    tools=[final_answer], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()