import random
import time
from datasets import load_dataset, Dataset, DatasetDict
from huggingface_hub import HfApi, HfFolder
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# -------------------------------------
# CONFIGURAÇÕES DO MODELO LOCAL
# -------------------------------------
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
DATASET_REPO = "caarleexx/AM1"
LOCAL_DATA_FILE = "data.json"

print("[LOG] Carregando modelo...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="cpu"
)
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=200,
    temperature=0.9
)

# -------------------------------------
# FUNÇÃO: gerar pergunta + N opções
# -------------------------------------
def gerar_pergunta_e_opcoes():
    N = random.randint(1, 3)  # número aleatório de alternativas
    prompt = (
        f"Crie uma pergunta curta. "
        f"A pergunta deve ter exatamente {N} respostas possíveis, todas bem diferentes entre si."
    )

    print(f"[LOG] Gerando pergunta com N={N} opções...")
    pergunta = generator(prompt)[0]["generated_text"]

    return pergunta, N

# -------------------------------------
# FUNÇÃO: gerar resposta principal
# -------------------------------------
def gerar_resposta(pergunta):
    print(f"[LOG] Gerando resposta principal...")
    resp = generator(f"{pergunta}\nResponda de forma curta e direta.")[0]["generated_text"]
    return resp

# -------------------------------------
# FUNÇÃO: gerar alternativas divergentes
# -------------------------------------
def gerar_alternativas(pergunta, N):
    prompt = (
        f"{pergunta}\n"
        f"Liste {N} respostas possíveis, todas diferentes entre si, "
        f"sem justificativas. a rssposta deve conter Apenas a lista."
    )

    print(f"[LOG] Gerando {N} alternativas divergentes...")
    texto = generator(prompt)[0]["generated_text"]

    alternativas = []
    for linha in texto.split("\n"):
        linha = linha.strip("-•1234567890. ")
        if len(linha) > 2:
            alternativas.append(linha)

    return alternativas[:N]

# -------------------------------------
# FUNÇÃO: calcular divergência
# -------------------------------------
def divergencia(resposta_principal, alternativas):
    print("[LOG] Calculando divergência...")
    divs = []
    base = resposta_principal.lower()

    for alt in alternativas:
        alt_l = alt.lower()
        diff = abs(len(base) - len(alt_l))
        divs.append(diff)

    return divs

# -------------------------------------
# FUNÇÃO: gerar explicação final
# -------------------------------------
def gerar_explicacao(pergunta):
    prompt = (
        f"Explique, de forma detalhada, por que a pergunta abaixo pode gerar respostas tão diferentes:\n"
        f"{pergunta} a sus respista deve ser clara para detalhsr conexao das a rsposra sem explicação ou justificativa"
    )

    print("[LOG] Gerando explicação...")
    return generator(prompt)[0]["generated_text"]

# -------------------------------------
# INICIAR OU CARREGAR DATASET
# -------------------------------------
def carregar_dataset():
    try:
        print("[LOG] Tentando carregar dataset remoto...")
        ds = load_dataset(DATASET_REPO)
        return ds
    except:
        print("[LOG] Criando dataset novo...")
        return DatasetDict({"train": Dataset.from_dict({"pergunta": [], "resposta": [], "alternativas": [], "divergencias": [], "explicacao": []})})

# -------------------------------------
# SALVAR NO HF DATASET A CADA LOOP
# -------------------------------------
def salvar_dataset(ds):
    print("[LOG] Salvando dataset no HuggingFace...")
    ds.push_to_hub(DATASET_REPO)

# -------------------------------------
# LOOP PRINCIPAL
# -------------------------------------
def loop_geracao(qtd=10):
    ds = carregar_dataset()

    for i in range(qtd):
        print("\n====================================")
        print(f"[LOG] ITERAÇÃO {i+1}")
        print("====================================")

        pergunta, N = gerar_pergunta_e_opcoes()
        resposta = gerar_resposta(pergunta)
        alternativas = gerar_alternativas(pergunta, N)
        divergencias = divergencia(resposta, alternativas)
        explicacao = gerar_explicacao(pergunta)

        novo_registro = {
            pergunta,
            resposta,
            alternativas,
            divergencias,
            explicacao
        }

        print("[LOG] Adicionando exemplo ao dataset...")
        ds["train"] = ds["train"].add_item(novo_registro)

        salvar_dataset(ds)

        time.sleep(1)

    print("[LOG] Concluído.")

# -------------------------------------
# EXECUTAR O GERADOR
# -------------------------------------
if __name__ == "__main__":
    loop_geracao(qtd=50)