import streamlit as st
import pandas as pd
import os
import json
import zipfile
import io

# Função para carregar anotações específicas (discursivas ou práticas)
def carregar_anotacoes(tipo, pasta="anotacoes"):
    anotacoes = {}
    if os.path.exists(pasta):
        for arquivo in os.listdir(pasta):
            if arquivo.endswith(f"_{tipo}.json"):
                usuario = arquivo.split("_")[0]
                with open(os.path.join(pasta, arquivo), "r", encoding="utf-8") as f:
                    anotacoes[usuario] = json.load(f)
    return anotacoes

# Função para calcular os totais de anotações
def calcular_totais(anotacoes_discursivas, anotacoes_praticas, df_discursivas, df_praticas):
    resultados = []

    total_questoes_discursivas = len(df_discursivas)
    total_criterios_discursivos = sum(len(q["formated_criteria"]) for _, q in df_discursivas.iterrows())
    total_provas_praticas = len(df_praticas)
    total_criterios_praticos = sum(len(q["formated_criteria"]) for _, q in df_praticas.iterrows())

    usuarios = set(anotacoes_discursivas.keys()).union(set(anotacoes_praticas.keys()))

    for usuario in usuarios:
        questoes_discursivas_completas = 0
        criterios_discursivos_anotados = 0
        provas_praticas_completas = 0
        criterios_praticos_anotados = 0

        if usuario in anotacoes_discursivas:
            for questao_id, criterios in anotacoes_discursivas[usuario].items():
                if questao_id in df_discursivas["id"].astype(str).values:
                    criterios_totais = len(df_discursivas[df_discursivas["id"] == int(questao_id)].iloc[0]["formated_criteria"])
                    criterios_anotados = len([k for k, v in criterios.items() if v in ["Sim", "Não"]])
                    if criterios_anotados == criterios_totais:
                        questoes_discursivas_completas += 1
                    criterios_discursivos_anotados += criterios_anotados

        if usuario in anotacoes_praticas:
            for questao_id, criterios in anotacoes_praticas[usuario].items():
                if questao_id in df_praticas["id"].astype(str).values:
                    criterios_totais = len(df_praticas[df_praticas["id"] == int(questao_id)].iloc[0]["formated_criteria"])
                    criterios_anotados = len([k for k, v in criterios.items() if v in ["Sim", "Não"]])
                    if criterios_anotados == criterios_totais:
                        provas_praticas_completas += 1
                    criterios_praticos_anotados += criterios_anotados

        resultados.append({
            "usuario": usuario,
            "questoes_discursivas_anotadas": f"{questoes_discursivas_completas}/{total_questoes_discursivas}",
            "total_criterios_discursivos_anotados": f"{criterios_discursivos_anotados}/{total_criterios_discursivos}",
            "provas_praticas_anotadas": f"{provas_praticas_completas}/{total_provas_praticas}",
            "criterios_praticos_anotados": f"{criterios_praticos_anotados}/{total_criterios_praticos}"
        })

    return resultados

# Função para calcular itens incompletos por usuário
def calcular_itens_incompletos(anotacoes_discursivas, anotacoes_praticas, df_discursivas, df_praticas):
    resultados_incompletos = []

    usuarios = set(anotacoes_discursivas.keys()).union(set(anotacoes_praticas.keys()))

    for usuario in usuarios:
        # Verificar discursivas
        if usuario in anotacoes_discursivas:
            for questao_id, criterios in anotacoes_discursivas[usuario].items():
                if questao_id in df_discursivas["id"].astype(str).values:
                    criterios_totais = df_discursivas[df_discursivas["id"] == int(questao_id)].iloc[0]["formated_criteria"]
                    criterios_anotados = [k for k, v in criterios.items() if v in ["Sim", "Não"]]
                    if len(criterios_anotados) < len(criterios_totais):
                        for criterio in criterios_totais:
                            if criterio not in criterios_anotados:
                                resultados_incompletos.append({
                                    "usuario": usuario,
                                    "item_pendente": f"D {questao_id} {criterio['letra']}-{criterio['parte']}"
                                })

        # Verificar práticas
        if usuario in anotacoes_praticas:
            for questao_id, criterios in anotacoes_praticas[usuario].items():
                if questao_id in df_praticas["id"].astype(str).values:
                    criterios_totais = df_praticas[df_praticas["id"] == int(questao_id)].iloc[0]["formated_criteria"]
                    criterios_anotados = [k for k, v in criterios.items() if v in ["Sim", "Não"]]
                    if len(criterios_anotados) < len(criterios_totais):
                        for criterio in criterios_totais:
                            if criterio not in criterios_anotados:
                                resultados_incompletos.append({
                                    "usuario": usuario,
                                    "item_pendente": f"P {questao_id} {criterio['numero']}-{criterio['parte']}"
                                })

    return resultados_incompletos

# Função principal do Streamlit
def run():
    #st.set_page_config(layout="wide")
    st.title("Dashboard de Anotações - Benchmark Jurídico OAB")

    # Carregar os DataFrames
    with st.spinner("Carregando arquivos..."):
        df_discursivas = pd.read_pickle("df_evaluated_discursive.pkl")
        df_praticas = pd.read_pickle("df_evaluated.pkl")
        anotacoes_discursivas = carregar_anotacoes("anotacoes_discursivas")
        anotacoes_praticas = carregar_anotacoes("anotacoes_praticas")

    # Calcular os totais
    resultados = calcular_totais(anotacoes_discursivas, anotacoes_praticas, df_discursivas, df_praticas)
    df_resultados = pd.DataFrame(resultados)

    # Exibir os dados em uma tabela
    st.markdown("### Resumo das Anotações")
    st.dataframe(df_resultados, use_container_width=True)

    # Calcular itens incompletos
    itens_incompletos = calcular_itens_incompletos(anotacoes_discursivas, anotacoes_praticas, df_discursivas, df_praticas)
    df_incompletos = pd.DataFrame(itens_incompletos)

    # Permitir filtragem por usuário
    usuarios_disponiveis = df_incompletos["usuario"].unique().tolist()
    usuario_selecionado = st.selectbox("Selecione um usuário:", options=["Todos"] + usuarios_disponiveis)

    if usuario_selecionado != "Todos":
        df_incompletos = df_incompletos[df_incompletos["usuario"] == usuario_selecionado]

    # Exibir os itens incompletos
    st.markdown("### Itens Incompletos por Usuário")
    st.dataframe(df_incompletos, use_container_width=True)

    # Criar um ZIP com os arquivos JSON
    folder_path = "anotacoes"
    zip_buffer = io.BytesIO()

    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
        for file_name in os.listdir(folder_path):
            if file_name.endswith(".json"):
                file_path = os.path.join(folder_path, file_name)
                zip_file.write(file_path, arcname=file_name)

    zip_buffer.seek(0)

    # Criar botão de download
    st.download_button(
        label="📥 Baixar Todos os JSONs",
        data=zip_buffer,
        file_name="anotacoes.zip",
        mime="application/zip"
    )