from src.agent import PlantUMLAgent
from src.documents import get_processed_documents, get_diagrams_names
from src.interface import create_interface
from src.preview import render_plantuml, sanitize_plantuml_code
import gradio as gr

from cachetools import LRUCache, cached

diagrams_names = get_diagrams_names()
docs_processed = get_processed_documents(diagrams_names)

cache = LRUCache(maxsize=100)

@cached(cache)
def respond(message, diagram_name, image_input, api_key):
    if not api_key:
        gr.Warning("Please provide a valid Mistral API key to generate diagrams.")
        return message, "", ""
    if not message:
        gr.Warning("Please provide a description for the UML diagram.")
        return "", "", ""
    
    agent = PlantUMLAgent(docs_processed, api_key)
    
    input = "Description:\n" + message.strip() + "\n\nDiagram Type: " + diagram_name.strip() + "\n\n"
    
    if image_input:
        image_analysis = agent.recognize_image(diagram_name, image_input)
        input += f"\n\nImage Analysis:\n{image_analysis}\n\n" if image_analysis else ""

    response = agent.predict(diagram_name, input)

    plantuml_code = sanitize_plantuml_code(response)    
    plantuml_diagram = render_plantuml(plantuml_code)
    
    return message, plantuml_code, plantuml_diagram

if __name__ == "__main__":
    cache.clear()
    demo = create_interface(diagrams_names, respond)
    demo.launch()