import gradio as gr
from huggingface_hub import InferenceClient
from datasets import load_dataset
import time

# Função para carregar o dataset `aicoder69/aicoder69` com tratamento de erros
def load_aicoder_dataset():
    try:
        print("Carregando o dataset...")
        dataset = load_dataset("aicoder69/aicoder69")
        print("Dataset carregado com sucesso!")
        return dataset
    except Exception as e:
        print(f"Erro ao carregar o dataset: {e}")
        return None

aicoder_dataset = load_aicoder_dataset()

# Função para recuperar um exemplo do dataset com segurança
def get_example_from_aicoder(dataset, index):
    if dataset and "train" in dataset:
        try:
            return dataset["train"][index]
        except IndexError:
            print("Índice fora do intervalo no dataset.")
            return {"text": "Nenhum exemplo disponível."}
    else:
        print("O dataset não foi carregado corretamente.")
        return {"text": "Dataset não disponível."}

# Inicializar o cliente de inferência com tratamento de erros
def initialize_client():
    try:
        print("Inicializando o cliente de inferência...")
        client = InferenceClient("unsloth/Qwen2.5-Coder-32B-Instruct")
        print("Cliente de inferência inicializado com sucesso!")
        return client
    except Exception as e:
        print(f"Erro ao inicializar o cliente de inferência: {e}")
        return None

client = initialize_client()

# Função de resposta do chatbot
def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    if not client:
        return "Erro: O cliente de inferência não foi inicializado."

    messages = [{"role": "system", "content": system_message}]

    # Adicionar interações históricas
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # Adicionar mensagem do usuário
    messages.append({"role": "user", "content": message})

    try:
        print("Enviando solicitação ao modelo...")
        # Ajuste o tempo limite e tente novamente caso haja falha
        retries = 3
        for attempt in range(retries):
            try:
                response = client.chat_completion(
                    messages,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                ).choices[0].message.content
                print("Resposta recebida com sucesso!")
                return response
            except Exception as e:
                print(f"Erro na tentativa {attempt + 1}/{retries}: {e}")
                if attempt < retries - 1:
                    print("Tentando novamente...")
                    time.sleep(2)  # Pausa entre tentativas
                else:
                    return f"Erro ao gerar resposta após {retries} tentativas."
    except Exception as e:
        print(f"Erro ao enviar solicitação: {e}")
        return "Ocorreu um erro ao gerar uma resposta."

# Exemplo: Recuperar uma entrada do dataset
example_data = get_example_from_aicoder(aicoder_dataset, 0)
print("Exemplo do dataset:", example_data)

# Interface Gradio
def launch_demo():
    try:
        demo = gr.ChatInterface(
            respond,
            additional_inputs=[
                gr.Textbox(value="Você é um chatbot amigável. Seu nome é Juninho.", label="Mensagem do sistema"),
                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Máximo de novos tokens"),
                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperatura"),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (amostragem núcleo)",
                ),
            ],
        )
        demo.launch()
    except Exception as e:
        print(f"Erro ao iniciar o aplicativo Gradio: {e}")

if __name__ == "__main__":
    launch_demo()