"""
Executor Direto na Instância Ativa (47729771 - 4x RTX 4090).
Executa Treinamento 1.1B com Gradient Checkpointing -> 3 Benchmarks -> Download -> Destruição.
"""

import os
import sys
import time
import subprocess

INSTANCE_ID = 47729771
SSH_HOST = "ssh8.vast.ai"
SSH_PORT = 19770
VAST_CLI = "/home/j/.local/bin/vastai"

ssh_opts = [
    "-o", "StrictHostKeyChecking=no",
    "-o", "UserKnownHostsFile=/dev/null",
    "-o", "ConnectTimeout=10",
    "-o", "ServerAliveInterval=15",
    "-p", str(SSH_PORT)
]
ssh_target = f"root@{SSH_HOST}"

def run_job():
    print("=" * 80)
    print(f"🌟 DISPARANDO TREINAMENTO 1.1B NA INSTÂNCIA ATIVA: {INSTANCE_ID} (4x RTX 4090)")
    print("=" * 80)

    try:
        # 1. Enviar Código Atualizado com Gradient Checkpointing
        print("\n▶ 1. Compactando e transmitindo código atualizado...")
        tarball = "/tmp/estrela_deploy_v2.tar.gz"
        subprocess.run([
            "tar", "-czf", tarball,
            "--exclude=node_modules",
            "--exclude=.git",
            "--exclude=estrela_core/from_scratch_llm/checkpoints_1b/*.safetensors",
            "estrela_core", "scripts"
        ], check=True)

        cat_proc = subprocess.Popen(["cat", tarball], stdout=subprocess.PIPE)
        subprocess.run(
            ["ssh"] + ssh_opts + [ssh_target, "mkdir -p /workspace && tar -xzf - -C /workspace"],
            stdin=cat_proc.stdout,
            check=True
        )
        cat_proc.stdout.close()
        print("✔ Código e dataset transmitidos!")

        # 2. Executar Treinamento Distribuído
        print("\n▶ 2. 🚀 INICIANDO TREINAMENTO DISTRIBUÍDO (4x RTX 4090 - bfloat16 + Gradient Checkpointing)...")
        train_cmd = """
cd /workspace && \
torchrun --nproc_per_node=4 estrela_core/from_scratch_llm/train_distributed_1b.py 2>&1 | tee /workspace/training.log
"""
        t0 = time.time()
        res = subprocess.run(["ssh"] + ssh_opts + [ssh_target, train_cmd], check=True)
        elapsed = time.time() - t0
        print(f"✔ Treinamento concluído com sucesso em {elapsed/60:.2f} minutos!")

        # 3. Executar Benchmarks Remotos
        print("\n▶ 3. 🧪 EXECUTANDO SUÍTE DE 3 BENCHMARKS NA GPU REMOTA...")
        bench_cmd = """
cd /workspace && \
python3 estrela_core/from_scratch_llm/benchmark_model.py \
  /workspace/estrela_core/from_scratch_llm/checkpoints_1b/estrelarosa_1b_epoch_3.safetensors \
  /workspace/estrela_core/from_scratch_llm/checkpoints_1b/tokenizer/tokenizer.json 2>&1 | tee /workspace/benchmark.log
"""
        subprocess.run(["ssh"] + ssh_opts + [ssh_target, bench_cmd], check=True)

        # 4. Baixar Pesos e Logs
        print("\n▶ 4. 📥 BAIXANDO PESOS FINAIS, TOKENIZER E BENCHMARKS...")
        local_ckpt_dir = "estrela_core/from_scratch_llm/checkpoints_1b"
        os.makedirs(local_ckpt_dir, exist_ok=True)

        download_remote = """
cd /workspace && tar -czf - \
  estrela_core/from_scratch_llm/checkpoints_1b/estrelarosa_1b_epoch_3.safetensors \
  estrela_core/from_scratch_llm/checkpoints_1b/tokenizer \
  benchmark_results.json training.log benchmark.log 2>/dev/null
"""
        dl_proc = subprocess.Popen(["ssh"] + ssh_opts + [ssh_target, download_remote], stdout=subprocess.PIPE)
        subprocess.run(["tar", "-xzf", "-", "-C", "."], stdin=dl_proc.stdout, check=True)
        dl_proc.stdout.close()
        print("✔ Todos os arquivos baixados com sucesso para o disco local!")

    finally:
        # 5. Destruição com confirmação automática
        print(f"\n▶ 5. 🛑 DESTRUINDO INSTÂNCIA {INSTANCE_ID} NO VAST.AI...")
        destroy_proc = subprocess.Popen(["echo", "y"], stdout=subprocess.PIPE)
        subprocess.run([VAST_CLI, 'destroy', 'instance', str(INSTANCE_ID)], stdin=destroy_proc.stdout)
        destroy_proc.stdout.close()
        print("✔ Instância destruída.")

    print("\n" + "=" * 80)
    print("🎉 OPERAÇÃO 100% CONCLUÍDA COM SUCESSO!")
    print("=" * 80)

if __name__ == "__main__":
    run_job()