import os
import logging
import json
import yaml
import base64
import tempfile

from smolagents import CodeAgent, HfApiModel, load_tool, tool
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
from plantuml import PlantUML
from smolagents.agent_types import AgentImage
from PIL import Image

# ------------------------------
# UML Generation Tool
# ------------------------------
@tool
def generate_uml_class_diagram(description: str, extra_details: str = "") -> object:
    """
    Generates a UML CLASS diagram image (in PlantUML syntax) from a system description.
    The tool instructs the LLM to assess the provided system description and,
    if necessary, add reasonable assumptions to generate complete UML code.
    
    It then enters a loop, re-calling the model until the output is clean—that is,
    it starts with '@startuml' and ends with '@enduml' with no extra text.
    
    Once a clean code block is obtained, it sends the code to a PlantUML server to produce
    a PNG diagram, saves the image to a temporary file, opens it with PIL, and returns the image
    wrapped as an AgentImage so that the Gradio UI displays it directly.
    
    Args:
        description: A natural language description (or topic) of the system.
        extra_details: (Optional) Additional details to supplement the description.
    
    Returns:
        An AgentImage object wrapping the generated UML diagram image if successful,
        otherwise a string with an error message.
    """
    # Combine description and extra details.
    full_description = description.strip()
    if extra_details.strip():
        full_description += "\n" + extra_details.strip()
    
    # Build the initial prompt.
    prompt = (
        "You are an expert in UML class diagramming. Given the following system description, "
        "generate complete PlantUML code for a UML class diagram that represents the system. "
        "Do not include any extra commentary—only output the UML code, starting with '@startuml' "
        "and ending with '@enduml'.\n\n"
        f"System Description:\n{full_description}\n\nPlantUML Code:\n@startuml"
    )
    
    max_attempts = 5
    attempt = 0
    plantuml_code = ""
    messages = [{"role": "user", "content": prompt}]
    
    # Loop until the output is clean.
    while attempt < max_attempts:
        try:
            response_obj = model(messages, stop_sequences=["@enduml"], max_tokens=2096)
            raw_output = response_obj.content if hasattr(response_obj, "content") else response_obj
        except Exception as e:
            logging.error("Error calling model: " + str(e))
            raw_output = ""
        
        start_idx = raw_output.find("@startuml")
        end_idx = raw_output.rfind("@enduml")
        if start_idx != -1 and end_idx != -1:
            candidate_code = raw_output[start_idx:end_idx + len("@enduml")]
        else:
            candidate_code = raw_output
        
        # Check for any extra text.
        before = raw_output[:start_idx].strip() if start_idx != -1 else ""
        after = raw_output[end_idx + len("@enduml"):].strip() if end_idx != -1 else ""
        if before == "" and after == "":
            plantuml_code = candidate_code
            break
        else:
            prompt = (
                "Please provide only the UML code in PlantUML syntax with no extra text. "
                "Your answer must start with '@startuml' and end with '@enduml'.\n\n"
                f"System Description:\n{full_description}\n\nPlantUML Code:\n@startuml"
            )
            messages = [{"role": "user", "content": prompt}]
        attempt += 1

    if not plantuml_code:
        plantuml_code = candidate_code

    # Process the generated PlantUML code to generate a PNG diagram.
    pl = PlantUML(url="http://www.plantuml.com/plantuml/img/")
    try:
        diagram_data = pl.processes(plantuml_code)
    except Exception as e:
        logging.error("Error generating diagram: " + str(e))
        diagram_data = None

    if diagram_data:
        try:
            # Save the diagram data to a temporary PNG file.
            tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png", mode="wb")
            tmp_file.write(diagram_data)
            tmp_file.close()
            # Open the image with PIL.
            pil_image = Image.open(tmp_file.name)
            # Wrap the image as an AgentImage and store the image object.
            agent_image = AgentImage(pil_image)
            agent_image.image = pil_image  # Ensure the image attribute is set
            return agent_image
        except Exception as e:
            logging.error("Error saving diagram to file: " + str(e))
            return f"Error saving diagram to file: {str(e)}"
    else:
        return f"Error generating diagram. Generated PlantUML code was:\n{plantuml_code}"

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

final_answer = FinalAnswerTool()

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, generate_uml_class_diagram],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()