# app.py

import streamlit as st
import requests
import json
from transformers import pipeline
from questions import questions

st.title("Classificação de Comentários do CEU")
st.write("Este aplicativo baixa respostas da API, classifica os comentários e exibe os resultados.")

@st.cache(allow_output_mutation=True)
def get_questions_map(questions_list):
    return {str(q["id"]): q["question"] for q in questions_list}

@st.cache(allow_output_mutation=True)
def load_classifier():
    classifier = pipeline(
        task="zero-shot-classification",
        model="joeddav/xlm-roberta-large-xnli"  # Modelo multilíngue
    )
    return classifier

def fetch_data(api_url):
    response = requests.get(api_url)
    if response.status_code == 200:
        return response.json()
    else:
        st.error(f"Erro ao acessar a API: {response.status_code}")
        return None

def process_data(data, questions_map, classifier, candidate_labels):
    associacoes = []
    classified_comments = []
    
    for registro in data:
        token = registro.get("student_info", {}).get("token", "Sem Token")
        answers = registro.get("answers", {})
        comments = registro.get("comments", {})
        
        for pergunta_id_str, resposta in answers.items():
            pergunta_texto = questions_map.get(pergunta_id_str, f"Pergunta {pergunta_id_str}")
            comentario = comments.get(pergunta_id_str, "Sem comentário")
            
            associacoes.append({
                "token": token,
                "perguntaId": pergunta_id_str,
                "perguntaTexto": pergunta_texto,
                "resposta": resposta,
                "comentario": comentario
            })
            
            # Classificação do comentário, se existir
            if comentario and comentario.strip():
                text_for_model = f"Pergunta: {pergunta_texto}. Comentário: {comentario}"
                classification_result = classifier(
                    text_for_model,
                    candidate_labels,
                    multi_label=False
                )
                
                best_label = classification_result["labels"][0]
                best_score = classification_result["scores"][0]
                
                classified_comments.append({
                    "token": token,
                    "perguntaId": pergunta_id_str,
                    "perguntaTexto": pergunta_texto,
                    "resposta": resposta,
                    "comentario": comentario,
                    "classificacao": best_label,
                    "confiança": round(best_score, 4)
                })
    
    return associacoes, classified_comments

def main():
    api_url = st.text_input("URL da API", "https://proced.datasavvy.com.br/api/responses")
    if st.button("Processar Dados"):
        data = fetch_data(api_url)
        if data:
            with st.spinner("Carregando modelo de classificação..."):
                classifier = load_classifier()
            questions_map = get_questions_map(questions)
            candidate_labels = [
                "sugestões de melhoria",
                "elogios",
                "reclamações",
                "solicitações de serviços ou infraestrutura"
            ]
            with st.spinner("Processando e classificando comentários..."):
                associacoes, classified_comments = process_data(data, questions_map, classifier, candidate_labels)
            
            st.success("Processamento concluído!")
            
            st.header("Associações de Perguntas, Respostas e Comentários")
            st.dataframe(associacoes)
            
            st.header("Comentários Classificados")
            st.dataframe(classified_comments)
            
            # Opção para baixar os resultados
            st.subheader("Baixar Resultados")
            associacoes_json = json.dumps(associacoes, indent=4, ensure_ascii=False)
            classified_comments_json = json.dumps(classified_comments, indent=4, ensure_ascii=False)
            
            st.download_button(
                label="Baixar Associações (JSON)",
                data=associacoes_json,
                file_name="associacoes.json",
                mime="application/json"
            )
            
            st.download_button(
                label="Baixar Classificações (JSON)",
                data=classified_comments_json,
                file_name="classified_comments.json",
                mime="application/json"
            )

if __name__ == "__main__":
    main()