import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq,
)
from datasets import load_dataset
import os
from huggingface_hub import login, HfApi

# ✅ Configuração Inicial
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"  # Melhora alocação de memória GPUs

# ✅ Autenticação no Hugging Face
HF_API_KEY = os.getenv("HF_API_KEY")
if not HF_API_KEY:
    raise ValueError("❌ ERRO: Token Hugging Face não encontrado no ambiente.")

login(HF_API_KEY)

# ✅ Definição de Variáveis
DATASET_NAME = "rwayz/reasoning"  # Dataset personalizado
MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct"  # Modelo base
NEW_MODEL_NAME = "APPONTE/reasoning-llama-8b"  # Novo modelo treinado

# ✅ Carregamento do Dataset
print(f"🔄 Carregando dataset: {DATASET_NAME}...")
dataset = load_dataset(DATASET_NAME)
print("✅ Dataset carregado com sucesso!")

# ✅ Exibir estrutura do dataset
print("📌 Colunas do dataset:", dataset["train"].column_names)

# ✅ Função para filtrar exemplos inválidos
def filter_invalid_examples(example):
    fields = [
        "IdCandidato", "Id", "Tags", "Url", "Url Perfil", "Cadastrado", "Nome", "Sexo",
        "Data de Nascimento", "Documento Pessoal", "Cidade", "E-mail", "Telefone", "Etapa atual",
        "Status da etapa", "Status na vaga", "Plataforma", "Objetivo profissional", "Tipo formacao", "Nome formacao", "Instituicao formacao", "Inicio formacao",
        "Fim formacao", "Conclusao formacao", "Empresa experiencia", "Cargo experiencia", "Inicio experiencia",
        "Fim experiencia", "Descricao atividades", "Tempo de experiencias formatado",
        "Tempo formatado calculado", "Tipo idioma", "Nivel idioma", "Criado em", "Atualizado em", "Pontuacao Final", "Classificacao Final", "Top", "Aderencia ao Perfil",
        "Status Profissional", "porcentagem_escolaridade", "porcentagem_posgraduacao", "porcentagem_idioma",
        "porcentagem_experiencias", "porcentagem_experiencia_descricao", "Inconsistencia Experiencia",
        "porcentagem_experiencia_desejada", "porcentagem_experiencia_trajetoria", "porcentagem_experiencia_instituicao_ou_setores",
        "porcentagem_cargos_e_funcaos", "porcentagem_certificacao", "porcentagem_estabilidade", "porcentagem_cargos_recentes"
    ]
    return any(bool(str(example.get(field, '')).strip()) for field in fields)

# 🔄 Remover exemplos inválidos
print("🔄 Removendo exemplos inválidos...")
dataset = dataset.filter(filter_invalid_examples)
print(f"✅ Dataset filtrado! {len(dataset['train'])} exemplos restantes.")

# ✅ Função para tratar dados nulos e tokenizar corretamente
def tokenize_function(examples):
    inputs = ""
    targets = ""

    # 🔹 Coletar e tratar todos os campos, convertendo nulos em strings vazias
    fields = [
        "IdCandidato", "Id", "Tags", "Url", "Url Perfil", "Cadastrado", "Nome", "Sexo",
        "Data de Nascimento", "Documento Pessoal", "Cidade", "E-mail", "Telefone", "Etapa atual",
        "Status da etapa", "Status na vaga", "Plataforma", "Objetivo profissional", "Tipo formacao", "Nome formacao", "Instituicao formacao", "Inicio formacao",
        "Fim formacao", "Conclusao formacao", "Empresa experiencia", "Cargo experiencia", "Inicio experiencia",
        "Fim experiencia", "Descricao atividades", "Tempo de experiencias formatado",
        "Tempo formatado calculado", "Tipo idioma", "Nivel idioma", "Criado em", "Atualizado em", "Pontuacao Final", "Classificacao Final", "Top", "Aderencia ao Perfil",
        "Status Profissional", "porcentagem_escolaridade", "porcentagem_posgraduacao", "porcentagem_idioma",
        "porcentagem_experiencias", "porcentagem_experiencia_descricao", "Inconsistencia Experiencia",
        "porcentagem_experiencia_desejada", "porcentagem_experiencia_trajetoria", "porcentagem_experiencia_instituicao_ou_setores",
        "porcentagem_cargos_e_funcaos", "porcentagem_certificacao", "porcentagem_estabilidade", "porcentagem_cargos_recentes"
    ]

    # 🔹 Criar uma string formatada com todas as informações relevantes
    inputs = "\n".join([f"{field}: {str(examples.get(field, '') or '').strip()}" for field in fields])
    
    # 🔹 Definir como alvo a pontuação final, classificação e aderência ao perfil
    targets = f"Pontuacao Final: {str(examples.get('Pontuacao Final', '')).strip()} | Classificacao Final: {str(examples.get('Classificacao Final', '')).strip()} | Aderencia ao Perfil: {str(examples.get('Aderencia ao Perfil', '')).strip()}"

    # 🔹 Tokenizar input_ids e labels
    model_inputs = tokenizer(inputs, padding="max_length", truncation=True, max_length=512)
    labels = tokenizer(targets, padding="max_length", truncation=True, max_length=512)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# 🔹 Aplicar tokenização corretamente e adicionar labels
print("🔄 Tokenizando dataset...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_API_KEY)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    token=HF_API_KEY
)

# ✅ Aplicar a tokenização ao novo dataset
tokenized_datasets = dataset.map(tokenize_function, batched=False, remove_columns=dataset["train"].column_names)
print("✅ Dataset tokenizado!")

# ✅ Criar Data Collator para empacotar os dados
from transformers import DataCollatorForSeq2Seq
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)

# ✅ Ajuste de Hiperparâmetros para LOGS DETALHADOS
training_args = TrainingArguments(
    output_dir="./results",
    eval_strategy="steps",
    eval_steps=1000,  # 🔹 Avaliação frequente
    save_strategy="steps",
    save_steps=2500,  # 🔹 Salvar checkpoints
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=4,
    num_train_epochs=10,  # 🔹 Mais épocas para um treinamento robusto
    weight_decay=0.01,
    logging_dir="./logs",
    logging_strategy="steps",
    logging_steps=50,  # 🔹 Logs frequentes
    save_total_limit=2,
    push_to_hub=True,
    hub_model_id=NEW_MODEL_NAME,
    hub_token=HF_API_KEY,
    gradient_checkpointing=True,
    bf16=True,
    learning_rate=1e-6,  # 🔹 Aprendizado mais lento
    max_grad_norm=0.2,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    optim="adamw_torch"
)

# ✅ Configurar Trainer para Fine-Tuning
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["train"],
    data_collator=data_collator,
)

# ✅ Iniciar Treinamento
print("🚀 Iniciando treinamento do modelo...")
try:
    trainer.train()
    print("✅ Treinamento concluído!")
except RuntimeError as e:
    print(f"❌ ERRO NO TREINAMENTO: {str(e)}")

# ✅ Salvar o Modelo Treinado Localmente
print("💾 Salvando modelo localmente...")
model.save_pretrained("./results")
tokenizer.save_pretrained("./results")
print("✅ Modelo salvo localmente em './results'!")

# ✅ Upload do modelo treinado para Hugging Face
print(f"🔄 Enviando modelo para o Hugging Face: {NEW_MODEL_NAME}...")
api = HfApi()
api.upload_folder(
    folder_path="./results",
    repo_id=NEW_MODEL_NAME,
    repo_type="model",
    token=HF_API_KEY
)

print(f"✅ Modelo salvo no Hugging Face: {NEW_MODEL_NAME}!")
print(f"🔗 Link do modelo: https://huggingface.co/rwayz/{NEW_MODEL_NAME}")