#!/usr/bin/env bash
# Treina um LoRA do Llama-3.1-8B-Instruct com os atendimentos exportados (estágio 2),
# usando MLX no Apple Silicon. Ajustado para 16 GB de RAM (M4): modelo 4-bit, batch 1.
#
# Pré-requisito: dataset exportado (backend: node scripts/exportarDatasetTreinamento.js)
# Uso: ./treinar_lora_mlx.sh [epocas]   (default 3)
set -euo pipefail
cd "$(dirname "$0")"

MODEL="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
DATA_DIR="data"
ADAPTER_DIR="adapters"
EPOCAS="${1:-3}"

if [[ ! -f "$DATA_DIR/train.jsonl" ]]; then
  echo "ERRO: $DATA_DIR/train.jsonl não existe. Rode antes: cd .. && node scripts/exportarDatasetTreinamento.js" >&2
  exit 1
fi

if [[ ! -d .venv ]]; then
  python3 -m venv .venv
fi
source .venv/bin/activate
pip -q install --upgrade mlx-lm

echo "Filtrando amostras por tamanho real em tokens (evita picos de memória na validação)"
python3 filtrar_por_tokens.py --data-dir "$DATA_DIR"

N_TRAIN=$(wc -l < "$DATA_DIR/train.jsonl" | tr -d ' ')
ITERS_TOTAL=$(( N_TRAIN * EPOCAS ))
if (( ITERS_TOTAL < 100 )); then ITERS_TOTAL=100; fi

# retomada: interrupções por falta de memória (rodando junto com outros processos
# pesados) não custam o progresso já feito. O mlx_lm.lora sempre reinicia a
# contagem de iteração em 1 a cada chamada (mesmo com --resume-adapter-file), então
# guardamos o total acumulado em $ADAPTER_DIR/.progresso.
#
# Importante: quando o processo é morto de fora (harness/SO por falta de memória),
# a árvore de processos inteira morre de uma vez — o bash NUNCA chega a rodar
# código depois do `python3 ...`, então atualizar o .progresso só no final não
# funciona. Por isso o progresso é escrito ao vivo, linha a linha, toda vez que o
# mlx_lm reporta ter salvo um checkpoint (--save-every), não só quando o treino
# termina. Os arquivos numerados (0000100_adapters.safetensors etc.) podem colidir
# entre retomadas, já que o contador local sempre reinicia em 1 — isso é só
# histórico/debug; o que importa pra retomar corretamente é o par
# adapters.safetensors (sempre a versão mais recente) + .progresso (iteração
# global correspondente).
PROGRESSO_FILE="$ADAPTER_DIR/.progresso"
ITERS_FEITAS=0
[[ -f "$PROGRESSO_FILE" ]] && ITERS_FEITAS=$(cat "$PROGRESSO_FILE")

RESUME_FLAGS=()
if (( ITERS_FEITAS > 0 )) && [[ -f "$ADAPTER_DIR/adapters.safetensors" ]]; then
  RESUME_FLAGS=(--resume-adapter-file "$ADAPTER_DIR/adapters.safetensors")
  echo "Retomando de $ITERS_FEITAS/$ITERS_TOTAL iterações já feitas"
fi

ITERS=$(( ITERS_TOTAL - ITERS_FEITAS ))
if (( ITERS < 1 )); then
  echo "Já foram feitas $ITERS_FEITAS/$ITERS_TOTAL iterações — nada a treinar. Apague $ADAPTER_DIR/ pra recomeçar do zero." >&2
  exit 0
fi

echo "Treinando: $N_TRAIN amostras, $EPOCAS épocas => $ITERS_TOTAL iterações no total ($ITERS_FEITAS já feitas, $ITERS restantes, batch 1)"

set +e
PYTHONUNBUFFERED=1 python3 run_lora_seguro.py \
  --model "$MODEL" \
  --train \
  --data "$DATA_DIR" \
  --adapter-path "$ADAPTER_DIR" \
  "${RESUME_FLAGS[@]}" \
  --batch-size 1 \
  --num-layers 16 \
  --iters "$ITERS" \
  --max-seq-length 2048 \
  --grad-checkpoint \
  --save-every 100 \
  --steps-per-eval 100 \
  --val-batches 10 2>&1 | while IFS= read -r linha; do
    echo "$linha"
    if [[ "$linha" == *"Saved adapter weights"* ]]; then
      local_iter=$(grep -oE '^Iter [0-9]+' <<<"$linha" | grep -oE '[0-9]+')
      [[ -n "$local_iter" ]] && echo "$(( ITERS_FEITAS + local_iter ))" > "$PROGRESSO_FILE"
    fi
  done
TREINO_EXIT=${PIPESTATUS[0]}
set -e

if (( TREINO_EXIT == 0 )); then
  echo "$ITERS_TOTAL" > "$PROGRESSO_FILE"
else
  echo "Treino interrompido (exit $TREINO_EXIT). Progresso salvo: $(cat "$PROGRESSO_FILE" 2>/dev/null || echo "$ITERS_FEITAS")/$ITERS_TOTAL. Rode o script de novo pra continuar." >&2
  exit "$TREINO_EXIT"
fi

echo
echo "Adapter salvo em $ADAPTER_DIR/. Teste rápido:"
echo "  mlx_lm.generate --model $MODEL --adapter-path $ADAPTER_DIR --prompt 'Boa tarde, minha internet caiu'"
echo "Para publicar no Ollama: ./publicar_no_ollama.sh"