import gradio as gr
import re
import os
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_together import Together
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.document_loaders import TextLoader

token = os.environ['token']

reduce_template = """ A seguir está uma lista de resumos: {docs}
Primeiro, identifique as pessoas ou entidades envolvidas no processo para entender o contexto.
Em seguida, explique por que cada parte está envolvida.
Analise os resumos, destacando os principais pontos de cada um.
Com base nos pontos destacados, elabore um texto único que inclua todas as informações.
Não omita nenhum nome relatado.
Não invente nada que não esteja nos resumos.
Inclua as partes principais do texto jurídico: o relatório (resumo do caso), a fundamentação (razões da decisão) e a decisão (veredito final).
Use uma linguagem clara e simples para que todos possam entender, sem termos jurídicos complicados.
Isso não é uma explicação do caso, e sim uma reescrita com palavras simples e sem jargões jurídicos, faça com que um adolescente de 14 anos entenda.
Faça um Texto único Final como se estivesse explicando para um adolescente de 14 anos:"""

sem_resumir_template = """Com base nos documentos fornecidos: {docs}
Primeiramente, faça uma leitura atenta para entender a essência de cada documento. Identifique as pessoas, entidades e os principais elementos envolvidos, fornecendo uma visão clara do contexto.
Explique detalhadamente o papel de cada parte envolvida, oferecendo insights sobre suas motivações e a relevância de suas ações ou posições dentro do caso.
Analise os documentos com profundidade, destacando não apenas os pontos cruciais, mas também detalhes significativos que contribuem para a compreensão completa do caso.
Construa um texto abrangente que integre todas as informações relevantes. Evite omitir detalhes importantes que foram mencionados nos documentos, mantendo a fidelidade ao conteúdo original.
Assegure-se de incluir todos os aspectos fundamentais do texto jurídico: o relatório (que resume o caso), a fundamentação (que explica as razões da decisão) e a decisão (o veredito final), preservando a riqueza de detalhes.
Adote uma linguagem clara e acessível, evitando jargões jurídicos sem simplificar excessivamente o conteúdo, de modo que seja compreensível para um público amplo, incluindo adolescentes.
O objetivo é criar um texto informativo e rico em detalhes, que esclareça o caso de forma íntegra e educativa, como se estivesse apresentando a alguém com escolaridade básica, um adolescente de 14 anos não especializado na área jurídica."""

integro_template="""Trechos de um processo judicial: {docs}

Estamos mergulhando em alguns trechos de um processo judicial e queremos que você transforme essas informações em um texto claro e contínuo, adequado para alguém com educação básica. Para isso, siga os passos abaixo, mas ao invés de separar cada etapa, integre todas as informações em um único texto coeso. Use palavras simples e evite jargões jurídicos, mantendo a seriedade e estrutura do contexto legal. Aqui está o que você precisa fazer, resumido em um único processo:

- **Identifique as partes envolvidas no processo**, explicando de forma simples quem são e por que estão envolvidas.
- **Analise os trechos fornecidos**, destacando os pontos mais importantes de forma que fiquem claros para todos.
- **Elabore um texto único que inclua**:
  - Uma descrição clara do caso, incluindo o que aconteceu e quem está envolvido (Relatório).
  - As razões por trás da decisão tomada pelo juiz, simplificadas (Fundamentação).
  - O resultado final do caso, de maneira simplificada (Decisão).

O objetivo é criar um texto único e fluido que englobe todas essas informações de maneira compreensível para qualquer pessoa, mesmo aquelas com apenas a educação básica. Certifique-se de não adicionar informações novas e de não omitir nenhum nome mencionado nos trechos.
"""

resumo_template="""Sentenças do Caso: {docs}
Este documento contém uma série de sentenças relacionados a um caso jurídico. Primeiramente, vamos identificar as partes envolvidas e compreender suas motivações. Em seguida, sintetizaremos os aspectos cruciais de cada resumo, preservando os detalhes essenciais e o contexto.

Partes Envolvidas e Contexto:
Identifique claramente quem são as pessoas ou entidades principais do caso e por que elas estão envolvidas, fornecendo um breve contexto.

Análise e Síntese:
Consolide os resumos examinando os pontos principais e as evidências apresentadas. Destaque elementos chave que são comuns ou significativos em toda a documentação.

Texto Unificado:
Com base nos elementos destacados, redija um resumo que combine todas as informações essenciais. Este deve incluir: 
1. Relatório: um resumo do caso, apresentando os fatos principais.
2. Fundamentação: as razões principais que suportam a decisão, explicando o raciocínio por trás das conclusões.
3. Decisão: o veredito final ou as conclusões do caso.

Lembre-se de manter a objetividade, sem adicionar informações não presentes nos originais. Evite jargões jurídicos complexos, mas mantenha a precisão dos termos quando necessário para a compreensão do caso. O objetivo é criar um resumo claro, direto e informativo.

Instruções Finais:
Reveja o texto para garantir que seja compreensível e que não omita informações críticas. O resumo final deve permitir a qualquer pessoa entender os aspectos fundamentais do caso jurídico sem necessidade de consulta aos textos originais.
"""


model_used = "mistralai/Mixtral-8x7B-Instruct-v0.1"

def clean_text(text):
    text = text.replace('\xa0', ' ')
    text = re.sub(r'\s+', ' ', text)
    text = re.sub(r'\(#.*?\)', '', text)
    return text


def greet(uploaded_file, template_choice):
    loader = TextLoader(uploaded_file.name)
    docs = loader.load()

    for doc in docs:
      doc.page_content = clean_text(doc.page_content)

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=3000,
        chunk_overlap=300,
        separators=["\n\n", "\n", " ", ""]
    )

    if template_choice == "Template 2":
        prompt_template = PromptTemplate.from_template(integro_template)
        llm = Together(
          model=model_used,
          temperature=0.5,
          max_tokens=10000,
          top_k=60,
          together_api_key=token
        )
    elif template_choice == "Template 1 (menos resumido)":
        prompt_template = PromptTemplate.from_template(sem_resumir_template)
        llm = Together(
          model=model_used,
          temperature=0.5,
          max_tokens=10000,
          top_k=60,
          together_api_key=token
        )
    elif template_choice == "Template 3 (Resumo Apenas)":
        prompt_template = PromptTemplate.from_template(resumo_template)
        llm = Together(
          model=model_used,
          temperature=0,
          max_tokens=10000,
          top_k=60,
          together_api_key=token
        )
    else:
        prompt_template = PromptTemplate.from_template(reduce_template)
        llm = Together(
          model=model_used,
          temperature=0.5,
          max_tokens=10000,
          top_k=60,
          together_api_key=token
        )
    
    integro_chain = LLMChain(llm=llm, prompt=prompt_template)
    
    split_docs = text_splitter.split_documents(docs)

    return integro_chain.invoke(split_docs)['text']

iface = gr.Interface(
    fn=greet,
    inputs=[gr.File(file_types=['txt'], show_label=True),
           gr.Radio(["Template 1 (mais resumido)", "Template 1 (menos resumido)", "Template 2", "Template 3 (Resumo Apenas)"], label="Escolha o Template")],
    outputs=["text"],
    description="Esta ferramenta simplifica textos jurídicos para uma linguagem mais acessível. Escolha um template e carregue um arquivo .txt com o texto que deseja processar.",
)

iface.launch(share=True)