import gradio as gr
import json
import re
from huggingface_hub import InferenceClient

client = InferenceClient(
    "mistralai/Mistral-7B-Instruct-v0.1"
)

system_prompt = """[INST] <<SYS>>
You are a helpful, respectful and honest assistant. Always answer in spanish. Never answer in English. Always answer as
helpfully as possible, while being safe.

If a question does not make any sense, or is not factually coherent, explain
why instead of answering something not correct. If you don't know the answer
to a question, please don't share false information.

Your goal is to provide answers relating to a legal document, making classifications of arguments inside an allegation.<</SYS>>
"""

rag_text = """
"""

#format_text = """{
#    "tipo": "clasificación del argumento (Argumento tipo X)",
#    "descripcion": "Descripcion breve del argumento",
#    "texto_referencia": "texto original completo perteneciente a dicho argumento"
#}"""

format_text = """{
    "tipo": "clasificación del argumento (Argumento tipo X)",
    "descripcion": "Descripcion breve del argumento",
    "texto_referencia": "cita el fragmento del texto original letra a letra signo a signo en el cual has basado tu respuesta"
}"""

prompt_template_text = """
                          --- Lista Argumentos ---
    * Argumento tipo 49 (Desvirtuación de los hechos, el denunciado no está de acuerdo con el relato de los hechos)
    * Argumento tipo 304 (El denunciado solicita medios de prueba)
    * Argumento tipo 994 (El denunciado solicita fotografía como medio de prueba, desafiando el relato del agente denunciante)
    * Argumento tipo 1002 (Prescripción, el delito ha prescrito)
    * Argumento tipo 2014 (Principio de proporcionalidad, denunciado considera sanción excesiva)
    * Argumento tipo 2027 (Niega desobediencia de señales de tráfico de prohibición)
    * Argumento tipo 2002 (Denuncia no notificada en el momento en el que se formuló)
    * Argumento tipo 1001 (Defecto de forma en la denuncia, la sanción contiene algún error)
"""

recorte = """
CUARTA .- Que conforme al art.53.2 a) de la Ley 39/2015 de 1 de Octubre del Procedimiento
Administrativo Común de las Administraciones Públicas, SOLICITO la comunicación de la identidad
de la autoridad competente para imponerla sanción y de la norma que le atribuya tal competencia.
"""

def format_prompt(message):
  prompt = "<s>"
  prompt += f"[INST] {message} [/INST]"
  return prompt

def generate(prompt):

    generate_kwargs = dict(
        temperature=0.9,
        max_new_tokens=1024,
        top_p=0.95,
        repetition_penalty=1.0,
        do_sample=True,
        seed=42,
    )

    formatted_prompt = format_prompt(prompt)

    output = client.text_generation(formatted_prompt, **generate_kwargs)

    return output


def process_input(text, rag, prompt_template):

    #if not(rag):
    #    rag_text = " "
    #if not(prompt_template):
    #    prompt_template_text = " "
    
    prompt = f"""
    {system_prompt}

    Teniendo en cuenta que los argumentos se clasifican de la siguiente manera:
    {prompt_template_text}
    
    Partiendo del siguiente fragmento de texto:

    Fragmento:
    ---------------------------------------------------------------
      {text}
    --------------------------------------------------------------

    Identifica y clasifica el tipo de los argumentos expuestos en el anterior fragmento de texto siguiendo estos pasos cuidadosamente:
     - Primero: Determina cúantos argumentos se exponen en el fragmento.
     - Segundo: Determina el contenido de cada uno de los argumentos encontrados en el fragmento anterior.
     - Tercero: Clasifica cada uno de los argumentos encontrados en el fragmento anterior con el tipo de argumento que le corresponda.
     - Cuarto: Devuelve una lista de JSONs que contenga los argumentos ya clasificados con el siguiente formato:
        {format_text}


    Recuerda responder en español. Los argumentos que no se puedan clasificar o que *claramente* no encajen con ningún tipo predefinido serán clasificados como 'Desconocido'. Recuerda que los tipos de argumentos de representan con un número que puedes encontrar en la información anterior.
    Respira profundamente y piensa paso a paso.
    [/INST]


    """
    output = generate(prompt)
    
    return output,prompt


def create_interface():
    # Definimos los componentes de la interfaz
    input_text = gr.Textbox(label="Input")
    rag_checkbox = gr.Checkbox(label="RAG")
    prompt_template = gr.Checkbox(label="PromptTemplate")
    output_text = gr.Textbox(label="Output")
    classification_types_checkboxes = gr.CheckboxGroup(label="Clasificacion Tipo")

    def highlight_text(text, output, color="red"):
        # Encuentra las partes del texto que se encuentran en la variable output
        matches = re.findall(r"'texto_referencia'\s*:\s*\"(.*?)\"", output)
        # Resalta las partes del texto que se encuentran en la variable output en rojo
        for match in matches:
            if match in text:
                text = text.replace(match, f'<span style="color: {color};">{match}</span>')
        return text


    # Definimos la función que se ejecutará cuando se envíe la entrada
    def fn(text, rag, prompt_template):

        output = []
        output = process_input(text, rag, prompt_template)
        matches = re.findall(r'"tipo"\s*:\s*"(.*?)"', output[0])
        classification_types = matches
        classification_types_options = [(option, option) for option in classification_types]
        classification_types_checkboxes = gr.CheckboxGroup(label="Clasificacion Tipo", choices=classification_types_options,value=[option[0] for option in classification_types_options], interactive = True)

        # Resalta las partes del texto que se encuentran en la variable output en rojo
        highlighted_text = highlight_text(text, output[0], color="red")

        return highlighted_text, classification_types_checkboxes, output[0], output[1]
        
    examples = [
        [recorte, False, True]
    ]

    HF_TOKEN = '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'
    hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "MistralAI")

    # Creamos la interfaz
    iface = gr.Interface(
        fn=fn,
        inputs=[input_text, rag_checkbox, prompt_template],
        outputs=["html", classification_types_checkboxes, "text", "text"],
        examples=examples,
        allow_flagging="manual",
        flagging_options=["👍", "👎"],
        flagging_callback=hf_writer
    )

    return iface

iface = create_interface()
iface.launch()