import os
import streamlit as st
from groq import Groq
from io import BytesIO

# Set up the Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Function to request pattern generation from Groq API
def get_pattern_from_groq(body_measurements, garment_type):
    # Preparing the prompt for the LLM
    prompt = f"Generate a 2D {garment_type} pattern based on the following body measurements:\n{body_measurements}"
    
    # Making a request to Groq's API
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama3-8b-8192",
        stream=False
    )
    
    # Get the response content
    return chat_completion.choices[0].message.content

# Function to simulate creating a 2D garment pattern (for demo purposes)
def generate_garment_pattern(body_measurements, garment_type):
    # Here you would typically use the measurements to generate the 2D pattern.
    # We'll simulate this by returning a simple placeholder pattern.
    pattern = f"2D {garment_type} Pattern based on measurements:\n{body_measurements}"
    return pattern

# Main Streamlit app UI
def main():
    st.title("Garment Pattern Design Tool")
    
    # Select garment type
    garment_type = st.selectbox("Select Garment Type", ["Top Body Garment", "Bottom Garment"])
    
    # Collect body measurements based on garment type
    st.header(f"Enter Measurements for {garment_type}")
    
    body_measurements = {}
    
    if garment_type == "Top Body Garment":
        body_measurements['Chest'] = st.number_input('Chest (in cm)', min_value=0, step=1)
        body_measurements['Waist'] = st.number_input('Waist (in cm)', min_value=0, step=1)
        body_measurements['Neck'] = st.number_input('Neck (in cm)', min_value=0, step=1)
        body_measurements['Height'] = st.number_input('Height (in cm)', min_value=0, step=1)
    
    elif garment_type == "Bottom Garment":
        body_measurements['Waist'] = st.number_input('Waist (in cm)', min_value=0, step=1)
        body_measurements['Hip'] = st.number_input('Hip (in cm)', min_value=0, step=1)
        body_measurements['Inseam'] = st.number_input('Inseam (in cm)', min_value=0, step=1)
        body_measurements['Height'] = st.number_input('Height (in cm)', min_value=0, step=1)
    
    body_measurements_str = "\n".join([f"{key}: {value}" for key, value in body_measurements.items()])
    
    if st.button("Generate Pattern"):
        if all(value > 0 for value in body_measurements.values()):
            # Fetch pattern from Groq LLM
            st.write("Generating pattern...")
            pattern = get_pattern_from_groq(body_measurements_str, garment_type)
            
            # Simulate the 2D pattern generation (you can replace this with actual pattern generation logic)
            pattern_output = generate_garment_pattern(body_measurements_str, garment_type)
            
            # Display the result
            st.subheader(f"Generated 2D {garment_type} Pattern:")
            st.text(pattern_output)
            
            # Simulating file download
            buffer = BytesIO()
            buffer.write(pattern_output.encode())
            buffer.seek(0)
            
            st.download_button(
                label="Download Garment Pattern",
                data=buffer,
                file_name=f"{garment_type.lower().replace(' ', '_')}_pattern.txt",
                mime="text/plain"
            )
        else:
            st.error("Please enter valid measurements for all fields.")
    
if __name__ == "__main__":
    main()