import streamlit as st
from smolagents.agents import ToolCallingAgent
from smolagents import tool, LiteLLMModel
from typing import Optional
import cv2
import pytesseract
from PIL import Image
import io
import numpy as np
import base64

# Define the LiteLLMModel with OpenAI key
model = LiteLLMModel(model_id="gpt-4o", api_key="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")

@tool
def extract_components(image_data_base64: str) -> str:
    """
    Extract components from a web design image.

    Args:
        image_data_base64: The image data in base64 string format.

    Returns:
        A string describing the components found in the image.
    """
    image_data = base64.b64decode(image_data_base64)
    image = Image.open(io.BytesIO(image_data))
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2GRAY)
    components = pytesseract.image_to_string(gray)
    return components

@tool
def generate_code(components: str) -> str:
    """
    Generate code for the given components.

    Args:
        components: A string describing the components.

    Returns:
        The generated code for the components.
    """
    # This is a placeholder implementation. You can replace it with actual code generation logic.
    return f"Generated code for components: {components}"

# Define the ToolCallingAgent
agent = ToolCallingAgent(tools=[extract_components, generate_code], model=model)

# Streamlit app title
st.title("Web Design Component Extractor")

# File uploader for the web design image
uploaded_file = st.file_uploader("Upload a web design image", type=["png", "jpg", "jpeg"])

# Button to run the agent
if st.button("Extract and Generate Code"):
    if uploaded_file is not None:
        image_data = uploaded_file.read()
        image_data_base64 = base64.b64encode(image_data).decode('utf-8')
        components = agent.run(f"extract_components {image_data_base64}")
        code = agent.run(f"generate_code {components}")
        st.write("Extracted Components:", components)
        st.write("Generated Code:", code)
    else:
        st.write("Please upload an image.")