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

# Função para carregar usuários e senhas de um arquivo JSON
def carregar_usuarios(arquivo="users.json"):
    if os.path.exists(arquivo):
        with open(arquivo, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

# Função para carregar anotações de um arquivo específico para cada usuário
def carregar_anotacoes_usuario(usuario, pasta="anotacoes"):
    if not os.path.exists(pasta):
        os.makedirs(pasta)
    arquivo_usuario = os.path.join(pasta, f"{usuario}_anotacoes_discursivas.json")
    if os.path.exists(arquivo_usuario):
        with open(arquivo_usuario, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

# Função para salvar anotações de um usuário
def salvar_anotacoes_usuario(usuario, dados, pasta="anotacoes"):
    arquivo_usuario = os.path.join(pasta, f"{usuario}_anotacoes_discursivas.json")
    with open(arquivo_usuario, "w", encoding="utf-8") as f:
        json.dump(dados, f, ensure_ascii=False, indent=4)

# Função para autenticar o login
def autenticar(usuario, senha, usuarios):
    return usuarios.get(usuario) == senha

# Função para extrair questões do texto
def extrai_questoes(texto):
    enunciado = texto.split('A)')[0]
    opcao_A = texto.split('B)')[0].split('A)')[1]
    opcao_B = texto.split('B)')[1].split('Obs.')[0]
    disclaimer = texto.split('Obs.')[1]
    return {
        'enunciado': enunciado,
        'A': opcao_A,
        'B': opcao_B,
        'warning': disclaimer
    }

# Função principal da aplicação
def run():
    #st.set_page_config(layout="wide")

    # Tela de login
    usuarios = carregar_usuarios()
    if "usuario" not in st.session_state:
        st.session_state["usuario"] = None
    if "selected_criterio" not in st.session_state:
        st.session_state["selected_criterio"] = None

    if st.session_state["usuario"] is None:
        st.title("ICAIL - Benchmark jurídico OAB segunda fase")
        st.markdown("Questões discursivas")

        usuario = st.text_input("Usuário")
        senha = st.text_input("Senha", type="password")
        if st.button("Entrar"):
            if autenticar(usuario, senha, usuarios):
                st.session_state["usuario"] = usuario
                st.success(f"Bem-vindo, {usuario}!")
                st.rerun()
            else:
                st.error("Usuário ou senha inválidos.")
    else:
        usuario = st.session_state["usuario"]

        # Carregar anotações do usuário
        anotacoes_salvas = carregar_anotacoes_usuario(usuario)

        # Carregar os dados do arquivo pickle correspondente
        dados_pickle = "df_evaluated_discursive.pkl"
        df = pd.read_pickle(dados_pickle)

        # Configuração da barra lateral para escolher a questão
        st.sidebar.title(f"Usuário: {usuario}")
        questao_id = st.sidebar.selectbox("Selecione a questão", df["id"].tolist())

        # Filtrar o DataFrame para a questão selecionada
        questao = df[df["id"] == questao_id].iloc[0]

        # Inicializar o JSON para a questão selecionada, se não existir
        if str(questao_id) not in anotacoes_salvas:
            anotacoes_salvas[str(questao_id)] = {}

        # Layout principal
        col1, col2 = st.columns([2, 1])

        with col1:
            st.header("Leitura e Contexto")
            st.write(f"{questao['exam']} º exame da OAB, área Direito {questao['area']} - prova aplicada em {questao['exam date']}")

            tab1, tab2, tab3 = st.tabs(["Questão discursiva", "Gabarito", "Resposta do Candidato"])
            with tab1:
                st.markdown(f"{questao['question']}", unsafe_allow_html=True)
            with tab2:
                st.markdown(f"{questao['answer']}", unsafe_allow_html=True)
            with tab3:
                questoes = extrai_questoes(questao['question'])
                st.markdown(f"<strong>Enunciado:</strong><br>{questoes['enunciado']}<br>", unsafe_allow_html=True)
                st.markdown("<strong>Resposta do candidato:</strong>", unsafe_allow_html=True)
                respostas = json.loads(questao['model_answer'])
                for resposta in respostas:
                    for key, value in resposta.items():
                        st.markdown(f"<strong>{key}) {questoes[key]}</strong>", unsafe_allow_html=True)
                        st.markdown(f"{value}<br>", unsafe_allow_html=True)

        with col2:
            st.header("Lista de Critérios")
            for idx, criterio in enumerate(questao["formated_criteria"]):
                texto_botao = f"Critério {criterio['letra']} - {criterio['parte']}"
                if st.button(texto_botao, key=f"menu_criterio_{idx}"):
                    st.session_state["selected_criterio"] = idx
                    st.rerun()

        # Exibir o critério selecionado
        if st.session_state["selected_criterio"] is not None:
            criterio = questao["formated_criteria"][st.session_state["selected_criterio"]]
            st.subheader(f"Critério {criterio['letra']} - {criterio['parte']}")
            st.write(f"Gabarito: {criterio['gabarito']}")
            st.write(f"Critério: {criterio['criterio']} (Peso: {criterio['pontos']})")
            chave_anotacao = f"exame_{questao['exam']}_area_{questao['area']}_questao_numero_{questao['number']}_letra_{criterio['letra']}_parte_{criterio['parte']}"
            valor_salvo = anotacoes_salvas[str(questao_id)].get(chave_anotacao, "-")

            # Rádio para escolher a decisão
            atende = st.radio(
                f"Atende ao critério?",
                options=["-", "Sim", "Não"],
                index=["-", "Sim", "Não"].index(valor_salvo) if valor_salvo in ["-", "Sim", "Não"] else 0,
                key=f"criterio_{questao_id}_{st.session_state['selected_criterio']}"
            )

            # Atualizar o valor no JSON
            anotacoes_salvas[str(questao_id)][chave_anotacao] = atende

        # Botão para salvar as anotações
        if st.button("Salvar Anotação"):
            salvar_anotacoes_usuario(usuario, anotacoes_salvas)
            st.success("Anotação salva com sucesso!")
            st.write("Resultados da anotação:")
            st.json(anotacoes_salvas)

        # Opção para logout
        if st.sidebar.button("Logout"):
            st.session_state["usuario"] = None
            st.rerun()