import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForQuestionAnswering

def carregar_modelos():
    # Caminhos no Hugging Face Hub
    caminho_gpt = "mateuscasare/GPTNeo-Vade-Mecum-2023"
    caminho_bert = "mateuscasare/BERT-Vade-Mecum-2023"
    caminho_bert_juridico = "mateuscasare/modelo_juridico_pt"

    dispositivo = "cuda" if torch.cuda.is_available() else "cpu"

    # Carregar modelos do Hugging Face Hub
    tokenizer_gpt = AutoTokenizer.from_pretrained(caminho_gpt)
    model_gpt = AutoModelForCausalLM.from_pretrained(caminho_gpt).to(dispositivo)

    tokenizer_bert = AutoTokenizer.from_pretrained(caminho_bert)
    model_bert = AutoModelForQuestionAnswering.from_pretrained(caminho_bert).to(dispositivo)

    tokenizer_bert_juridico = AutoTokenizer.from_pretrained(caminho_bert_juridico)
    model_bert_juridico = AutoModelForQuestionAnswering.from_pretrained(caminho_bert_juridico).to(dispositivo)

    return tokenizer_gpt, model_gpt, tokenizer_bert, model_bert, tokenizer_bert_juridico, model_bert_juridico, dispositivo

def main():
    st.set_page_config(page_title="Assistente Jurídico", layout="wide")
    st.title("⚖️ Assistente Jurídico")

    # Carregar modelos
    tokenizer_gpt, model_gpt, tokenizer_bert, model_bert, tokenizer_bert_juridico, model_bert_juridico, dispositivo = carregar_modelos()

    # Interface principal
    tab1, tab2, tab3 = st.tabs(["Geração de Texto", "Perguntas e Respostas", "Perguntas Jurídicas Específicas"])

    with tab1:
        st.subheader("📝 Gerador de Textos Jurídicos")
        prompt = st.text_area("Descreva sua necessidade:", height=150, placeholder="Ex: Preciso de uma petição de habeas corpus...")

        col1, col2 = st.columns(2)
        with col1:
            max_length = st.slider("Tamanho da resposta", 100, 2000, 500)
            temperatura = st.slider("Criatividade", 0.0, 1.0, 0.7)

        if st.button("Gerar", type="primary"):
            with st.spinner("Gerando texto..."):
                inputs = tokenizer_gpt(prompt, return_tensors="pt").to(dispositivo)
                outputs = model_gpt.generate(**inputs, max_length=max_length, temperature=temperatura, do_sample=True)
                texto_gerado = tokenizer_gpt.decode(outputs[0], skip_special_tokens=True)
                st.text_area("Texto gerado:", texto_gerado, height=400)

    with tab2:
        st.subheader("❓ Consulta Jurídica (Modelo Genérico)")
        contexto = st.text_area("Contexto:", height=150, placeholder="Insira o contexto jurídico aqui...")
        pergunta = st.text_input("Sua pergunta:", placeholder="Ex: Qual o prazo para...?")

        if st.button("Responder (Modelo Genérico)", type="primary"):
            with st.spinner("Analisando..."):
                inputs = tokenizer_bert(pergunta, contexto, return_tensors="pt").to(dispositivo)
                outputs = model_bert(**inputs)
                answer_start = torch.argmax(outputs.start_logits)
                answer_end = torch.argmax(outputs.end_logits) + 1
                answer = tokenizer_bert.decode(inputs["input_ids"][0][answer_start:answer_end])
                st.success(f"Resposta: {answer}")

    with tab3:
        st.subheader("⚖️ Consulta Jurídica Específica")
        contexto_juridico = st.text_area("Contexto Jurídico:", height=150, placeholder="Insira o contexto jurídico específico aqui...")
        pergunta_juridica = st.text_input("Sua pergunta jurídica:", placeholder="Ex: Qual o prazo para recorrer?")

        if st.button("Responder (Modelo Jurídico)", type="primary"):
            with st.spinner("Analisando..."):
                inputs = tokenizer_bert_juridico(pergunta_juridica, contexto_juridico, return_tensors="pt").to(dispositivo)
                outputs = model_bert_juridico(**inputs)
                answer_start = torch.argmax(outputs.start_logits)
                answer_end = torch.argmax(outputs.end_logits) + 1
                answer = tokenizer_bert_juridico.decode(inputs["input_ids"][0][answer_start:answer_end])
                st.success(f"Resposta Jurídica: {answer}")

    with st.expander("ℹ️ Sobre os Modelos"):
        st.markdown("""
        Este assistente utiliza três modelos de IA:
        - **GPTNeo Vade-Mecum**: Para geração de textos jurídicos.
        - **BERT Vade-Mecum**: Para responder perguntas gerais sobre direito.
        - **Modelo Jurídico PT**: Treinado especificamente com textos jurídicos brasileiros para maior precisão.
        """)

if __name__ == "__main__":
    main()