import streamlit as st
import os
from dotenv import load_dotenv
from groq import Groq
import base64
from prompt import build_messages
from werkzeug.utils import secure_filename
from gerar_json import processar_imagem

# Carregar variáveis de ambiente
load_dotenv()
client = Groq(
    api_key=os.getenv("GROQ_API_KEY"),
)

# Função para codificar a imagem
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

# Configuração do Streamlit
st.title("Análise de Imagens de Vestuário")

# Formulário para entrada do usuário
with st.form(key="form_analise"):
    sexo = st.selectbox("Selecione o sexo:", ["Masculino", "Feminino"])
    ocasiao = st.selectbox("Selecione a ocasião:", ["Formal", "Casual", "Esporte", "Festa", "Praia"])
    imagem = st.file_uploader("Carregue a imagem:", type=["jpg", "jpeg", "png"])
    submit_button = st.form_submit_button(label="Enviar")

# Lógica de processamento ao enviar o formulário
if submit_button:
    if imagem is not None:
        # Garantir que o nome do arquivo seja seguro
        file_name = secure_filename(imagem.name)

        # Criar uma pasta temporária para armazenar o arquivo
        temp_dir = "temp_images"
        os.makedirs(temp_dir, exist_ok=True)
        file_path = os.path.join(temp_dir, file_name)

        # Salvar a imagem no disco
        with open(file_path, "wb") as f:
            f.write(imagem.getbuffer())

        # Codificar a imagem para base64
        base64_image = encode_image(file_path)

        # Processar a imagem
        deteccao = processar_imagem(file_path)[2]

        # Construir mensagens para os prompts
        messages_for_description, messages_for_analysis, format_answer = build_messages(deteccao, base64_image)

        try:
            # Primeira interação com o modelo
            chat_completion = client.chat.completions.create(
                messages=messages_for_description,
                model="llama-3.2-90b-vision-preview",
            )
            clothing_description = chat_completion.choices[0].message.content

            # Segunda interação com o modelo
            analysis_response = client.chat.completions.create(
                messages=messages_for_analysis(clothing_description, sexo, ocasiao),
                model="llama-3.2-11b-vision-preview",
            )
            resposta = analysis_response.choices[0].message.content

            final_answer = client.chat.completions.create(
                messages=format_answer(resposta),
                model="llama3-70b-8192"
            )
            resposta_final = final_answer.choices[0].message.content

            # Exibir o resultado
            st.success("Análise concluída com sucesso!")
            st.write("### Resultado:")
            st.write(resposta_final)

            os.remove(file_path)
        except Exception as e:
            st.error(f"Ocorreu um erro durante o processamento: {e}")
    else:
        st.warning("Por favor, carregue uma imagem para análise.")