import os
from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, tool
import datetime
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

# Tool to generate the RPA automation report
@tool
def generate_automation_report(case_description: str) -> str:
    """ Generates a structured Word report for RPA suitability analysis."""
    doc = Document()
    doc.add_heading('Automation Suitability & Requirements Report', 0)
    doc.add_heading('Case Description', 1)
    p = doc.add_paragraph(case_description)
    p.style = 'List Bullet'
    
    # Process Assessment Tool
    doc.add_heading('Process Assessment Tool (PAT)', 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'
    row = pat_table.add_row().cells
    row[0].text = 'Frequency of Execution'
    row[1].text = '5'
    row[2].text = 'Daily predictable patterns'
    
    # Requirements Template
    doc.add_heading('Business Requirements Template', 1)
    doc.add_paragraph('I. PROJECT SCOPE\nThe goal of this automation is to...')
    
    # PDD & SDD
    doc.add_heading('Process Definition Document (PDD)', 1)
    doc.add_heading('System Design Document (SDD)', 1)
    doc.add_paragraph('Technical Architecture\nThis section outlines the system architecture.')
    
    # UAT Plan
    doc.add_heading('User Acceptance Testing (UAT) Plan', 1)
    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'
    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'
    
    # Save report
    report_filename = 'Automation_Suitability_Report.docx'
    doc.save(report_filename)
    return f"Report generated: {report_filename} (update with specific details)."

# Model and tool setup
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None
)

final_answer = FinalAnswerTool()

# Configure agent with tools
agent = CodeAgent(
    model=model,
    tools=[final_answer, generate_automation_report],  # Include report tool
    max_steps=6,
    verbosity_level=1,
    prompt_templates=None  # Handle templates appropriately
)

# Launch Gradio UI
if __name__ == "__main__":
    GradioUI(agent).launch()