import os
from threading import Thread
from typing import Iterator

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

DESCRIPTION = """"\
# Bem-vindo ao Cícero Contratos!

Existo para ajudar você a analisar, gerar cláusulas contratuais e verificar se suas regras foram implementadas em um contrato. Veja exemplos abaixo de como você pode interagir comigo! Não sou advogado, sou um protótipo.
Neste momento eu recomendo você a usar duas habilidades principais:
A primeira é gerar cláusulas contratuais. Por exemplo, você quer algo para minimizar a responsabilidade no caso de descumprimento? Peça isso! 
A segunda é verificar se uma cláusula está de acordo com regras específicas. Por exemplo, sua empresa só assina contratos cujas cláusulas de eleição de foro indiquem o município de Ritápolis. Coloque a cláusula e a regra que você quer implementada e peça para o modelo adaptar. Economize tempo! 
É importante você ser direto. Use o verbo no imperativo: gere, produza, implemente, modifique. Peça por partes que a coisa funciona melhor! 
Se quiser falar com meu criador, seu e-mail é arthrod@umich.edu. 
"""

LICENSE = """
All rights reserved.
"""

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"


if torch.cuda.is_available():
    model_id = "arthrod/jaera"
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.use_default_system_prompt = False


@spaces.GPU
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    conversation = []
    if system_prompt:
        conversation.append({"role": "system", "content": system_prompt})
    for user, assistant in chat_history:
        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
    conversation.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=repetition_penalty,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)


chat_interface = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.Textbox(label="System prompt", lines=6),
        gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="Temperature",
            minimum=0.1,
            maximum=4.0,
            step=0.1,
            value=0.6,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.2,
        ),
    ],
    stop_btn=None,
    examples=[
        ["Na cl\u00e1usula a seguir, implemente a regra aplic\u00e1vel. Cl\u00e1usula: O Usu\u00e1rio concede \u00e0 [PARTE A] uma licen\u00e7a n\u00e3o exclusiva, mundial, perp\u00e9tua, irrevog\u00e1vel, livre de royalties, sublicenci\u00e1vel (atrav\u00e9s de m\u00faltiplos n\u00edveis) e transfer\u00edvel para usar, reproduzir, distribuir, preparar trabalhos derivados, exibir e executar o Conte\u00fado Gerado pelo Usu\u00e1rio em conex\u00e3o com o Servi\u00e7o e para incorporar o Conte\u00fado Gerado pelo Usu\u00e1rio em outros trabalhos, incluindo, mas n\u00e3o se limitando a, materiais de marketing e promo\u00e7\u00f5es. A [PARTE A] ter\u00e1 o direito, mas n\u00e3o a obriga\u00e7\u00e3o, de monitorar o Conte\u00fado Gerado pelo Usu\u00e1rio e poder\u00e1 remover qualquer Conte\u00fado Gerado pelo Usu\u00e1rio que viole esta licen\u00e7a ou que seja, a seu crit\u00e9rio exclusivo, considerado censur\u00e1vel. Regra: Defina os termos e condi\u00e7\u00f5es para o uso do conte\u00fado gerado pelo usu\u00e1rio, incluindo restri\u00e7\u00f5es e limita\u00e7\u00f5es de responsabilidade."],
        ["Na cl\u00e1usula a seguir, implemente a regra aplic\u00e1vel. Cl\u00e1usula: As Partes concordam em n\u00e3o fazer declara\u00e7\u00f5es p\u00fablicas desfavor\u00e1veis ou depreciativas sobre a outra Parte ou seus produtos, servi\u00e7os ou neg\u00f3cios. Esta obriga\u00e7\u00e3o permanecer\u00e1 em vigor durante o per\u00edodo de vig\u00eancia deste Contrato e por um per\u00edodo de [N\u00daMERO] anos ap\u00f3s o seu t\u00e9rmino. Regra: Evite linguagem difamat\u00f3ria ou prejudicial \u00e0 reputa\u00e7\u00e3o da outra parte."],
        ["Na cl\u00e1usula a seguir, implemente a regra aplic\u00e1vel. Cl\u00e1usula: A [PARTE A] poder\u00e1 exercer a op\u00e7\u00e3o de compra do [objeto da op\u00e7\u00e3o] nos termos deste instrumento mediante notifica\u00e7\u00e3o por escrito \u00e0 [PARTE B] at\u00e9 [DATA]. O exerc\u00edcio da op\u00e7\u00e3o dever\u00e1 ser acompanhado do pagamento do valor de [VALOR], corrigido monetariamente pela [\u00edndice de corre\u00e7\u00e3o monet\u00e1ria] desde a data de assinatura deste instrumento, referente ao pre\u00e7o de aquisi\u00e7\u00e3o do [objeto da op\u00e7\u00e3o]. Regra: Consequ\u00eancias: transfer\u00eancia de propriedade ap\u00f3s o exerc\u00edcio da op\u00e7\u00e3o de compra."],
        ["Na cl\u00e1usula a seguir, implemente a regra aplic\u00e1vel. Cl\u00e1usula: As Partes declaram e garantem que (i) possuem todo o poder e autoridade necess\u00e1rios para celebrar e cumprir este Contrato, (ii) est\u00e3o devidamente organizadas, v\u00e1lidas e existentes e possuem boa reputa\u00e7\u00e3o e cr\u00e9dito, (iii) t\u00eam o direito de celebrar e cumprir este Contrato sem violar quaisquer leis ou regulamentos aplic\u00e1veis ou quaisquer contratos ou acordos vinculativos, (iv) as informa\u00e7\u00f5es fornecidas \u00e0 outra Parte s\u00e3o verdadeiras, precisas e completas, e (v) n\u00e3o est\u00e3o sujeitas a quaisquer lit\u00edgios, a\u00e7\u00f5es ou processos judiciais que possam ter um efeito material adverso sobre sua capacidade de cumprir suas obriga\u00e7\u00f5es nos termos deste Contrato. As Partes concordam e se comprometem (i) a cumprir todas as leis, regulamentos, contratos e acordos aplic\u00e1veis relacionados ao objeto deste Contrato, (ii) a manter a outra Parte informada de quaisquer altera\u00e7\u00f5es em suas informa\u00e7\u00f5es fornecidas, (iii) a notificar a outra Parte imediatamente sobre quaisquer lit\u00edgios, a\u00e7\u00f5es ou processos judiciais que possam ter um efeito material adverso sobre sua capacidade de cumprir suas obriga\u00e7\u00f5es nos termos deste Contrato, (iv) a manter os padr\u00f5es de qualidade e desempenho conforme razoavelmente determinado pela outra Parte e (v) a n\u00e3o ceder ou transferir este Contrato ou quaisquer direitos ou obriga\u00e7\u00f5es aqui contidos sem o consentimento pr\u00e9vio por escrito da outra Parte. Regra: As representa\u00e7\u00f5es e garantias fornecem \u00e0s partes uma base para confiar nas alega\u00e7\u00f5es da outra parte sobre fatos importantes."],
        ["Redija um contrato abrangente que inclua disposi\u00e7\u00f5es para garantir um empr\u00e9stimo com garantia real na propriedade intelectual do mutu\u00e1rio, abrangendo direitos, recursos e a natureza cumulativa dessas disposi\u00e7\u00f5es"],
    ],
)

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
    chat_interface.render()
    gr.Markdown(LICENSE)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()