#!/usr/bin/env python3
"""
Khudi AI v2 — Safe Training Script
- Saves checkpoints to HF every 100 steps (auto-recovery)
- Health check: auto-stops if GPU idle for 10 min
- Resumes from latest checkpoint
- Saves final model with tokenizer
- Tests connectivity periodically
"""
import os
import sys
import time
import json
import shutil
import subprocess
import threading
from pathlib import Path

# ============================================================
# CONFIGURATION
# ============================================================
MODEL_NAME = "Qwen/Qwen3.5-9B-Instruct"
DATASET_NAME = "ZaoKing/khudi-ai-dataset"  # Will upload v2 to this repo
OUTPUT_DIR = "/workspace/khudi-v2-output"
HF_TOKEN = os.environ.get("HF_TOKEN", "")
CHECKPOINT_DIR = f"{OUTPUT_DIR}/checkpoints"

os.environ["HF_HOME"] = f"{OUTPUT_DIR}/hf_cache"
os.environ["HUGGINGFACE_HUB_TOKEN"] = HF_TOKEN
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

# ============================================================
# HEALTH CHECK THREAD
# ============================================================
class HealthChecker:
    def __init__(self, pid):
        self.pid = pid
        self.last_active = time.time()
        self.stop_flag = False

    def is_gpu_active(self):
        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"],
                capture_output=True, text=True, timeout=10
            )
            util = int(result.stdout.strip())
            return util > 0
        except:
            return False

    def start(self):
        def loop():
            while not self.stop_flag:
                time.sleep(60)  # Check every minute
                if self.is_gpu_active():
                    self.last_active = time.time()
                else:
                    idle_minutes = (time.time() - self.last_active) / 60
                    if idle_minutes > 10:
                        print(f"\n[HEALTH CHECK] GPU idle for {idle_minutes:.0f} min - training is stuck!")
                        print("[HEALTH CHECK] Saving checkpoint and exiting")
                        # Try to save and exit cleanly
                        os.system("pkill -f 'python.*train' 2>/dev/null")
                        time.sleep(5)
                        os._exit(2)  # Special exit code
                # Also check that we have network
                if not self.check_internet():
                    print(f"\n[HEALTH CHECK] No internet! Trying to recover...")
                    time.sleep(30)
        t = threading.Thread(target=loop, daemon=True)
        t.start()

    def check_internet(self):
        try:
            subprocess.run(["curl", "-sI", "--max-time", "5", "https://huggingface.co"],
                         capture_output=True, timeout=10)
            return True
        except:
            return False

    def stop(self):
        self.stop_flag = True


# ============================================================
# STEP 1: Install dependencies
# ============================================================
print("=" * 70)
print("🇵🇰 KHUDI AI v2 — SAFE TRAINING PIPELINE")
print("=" * 70)
print(f"Model: {MODEL_NAME}")
print(f"Dataset: {DATASET_NAME}")
print(f"Output: {OUTPUT_DIR}")
print()

print("[1/6] Installing dependencies...")
subprocess.run(["pip", "install", "-q",
                "transformers>=4.45.0",
                "datasets>=2.20.0",
                "peft>=0.10.0",
                "trl>=0.10.0",
                "accelerate>=0.30.0",
                "bitsandbytes>=0.43.0",
                "huggingface_hub>=0.24.0",
                "sentencepiece",
                "protobuf"], check=False)

# ============================================================
# STEP 2: Setup directories
# ============================================================
print("[2/6] Setting up directories...")
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
os.makedirs(f"{OUTPUT_DIR}/hf_cache", exist_ok=True)

# ============================================================
# STEP 3: Login to HF
# ============================================================
print("[3/6] Authenticating with Hugging Face...")
from huggingface_hub import HfApi, login
login(token=HF_TOKEN, add_to_git_credential=False)
api = HfApi()
print(f"  ✅ Logged in")

# ============================================================
# STEP 4: Download dataset
# ============================================================
print("[4/6] Downloading v2 dataset...")
from datasets import load_dataset, Dataset

try:
    # Try loading v2 dataset
    dataset = load_dataset(DATASET_NAME, split="train")
    print(f"  ✅ Loaded {len(dataset)} samples")
except Exception as e:
    print(f"  ⚠️  Could not load from HF: {e}")
    print("  Trying local fallback...")
    # Try local file
    local_files = ["/workspace/v2_data.jsonl", "/data/v2.jsonl", "/root/v2.jsonl"]
    for f in local_files:
        if os.path.exists(f):
            dataset = load_dataset("json", data_files=f, split="train")
            print(f"  ✅ Loaded {len(dataset)} from {f}")
            break
    else:
        print("  ❌ No dataset found!")
        sys.exit(1)

# Filter for shorter sequences (saves time + memory)
def length_filter(example):
    total_len = sum(len(m.get("content", "")) for m in example.get("messages", []))
    return 100 < total_len < 4000  # Min 100 chars, max 4000

dataset = dataset.filter(length_filter)
print(f"  ✅ After length filter: {len(dataset)} samples")

# ============================================================
# STEP 5: Check for existing checkpoint (resume capability)
# ============================================================
print("[5/6] Checking for resume checkpoints...")
resume_from = None
if os.path.exists(CHECKPOINT_DIR):
    checkpoints = [d for d in os.listdir(CHECKPOINT_DIR) if d.startswith("checkpoint-")]
    if checkpoints:
        checkpoints.sort(key=lambda x: int(x.split("-")[1]))
        latest = checkpoints[-1]
        resume_from = f"{CHECKPOINT_DIR}/{latest}"
        print(f"  ✅ Found checkpoint: {latest} - will resume from here")

# ============================================================
# STEP 6: Load model and train
# ============================================================
print("[6/6] Loading model and starting training...")
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

# BitsAndBytes config (4-bit for memory efficiency)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# Load model
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    attn_implementation="sdpa",  # More stable than flash_attn
)
model.config.use_cache = False
model = prepare_model_for_kbit_training(model)

# LoRA config
lora_config = LoraConfig(
    r=32,  # Higher rank for more capacity
    lora_alpha=64,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Format the dataset
def format_chat(example):
    msgs = example.get("messages", [])
    text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
    return {"text": text}

print("  Formatting dataset...")
dataset = dataset.map(format_chat, remove_columns=dataset.column_names)
print(f"  ✅ Formatted {len(dataset)} samples")

# Training arguments (with all safety settings)
training_args = SFTConfig(
    output_dir=CHECKPOINT_DIR,
    num_train_epochs=1,
    per_device_train_batch_size=2,  # Small for safety
    gradient_accumulation_steps=8,  # Effective batch = 16
    learning_rate=2e-4,
    max_length=1024,  # Shorter for speed
    logging_steps=10,
    save_steps=100,  # ⭐ CHECKPOINT EVERY 100 STEPS (was 500)
    save_total_limit=3,  # Keep only 3 most recent (save disk)
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    optim="paged_adamw_8bit",
    fp16=False,
    bf16=True,
    gradient_checkpointing=True,
    report_to="none",
    resume_from_checkpoint=resume_from is not None,
    push_to_hub=True,  # ⭐ AUTO-PUSH TO HF
    hub_model_id="ZaoKing/khudi-ai-v2",  # NEW v2 repo
    hub_token=HF_TOKEN,
    hub_strategy="checkpoint",  # Push at every checkpoint
    load_best_model_at_end=False,
    save_safetensors=True,
)

# Trainer
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
)

# Start health check
print("\n  🏥 Starting health check monitor...")
hc = HealthChecker(os.getpid())
hc.start()

# TRAIN
print("\n🚀 Starting training...")
print("=" * 70)
try:
    trainer.train(resume_from_checkpoint=resume_from)
    print("\n✅ Training completed successfully!")
except Exception as e:
    print(f"\n❌ Training error: {e}")
    print("Saving current state...")
    trainer.save_model(f"{OUTPUT_DIR}/emergency_save")
    api.upload_folder(
        folder_path=f"{OUTPUT_DIR}/emergency_save",
        repo_id="ZaoKing/khudi-ai-v2",
        commit_message=f"Emergency save after error: {str(e)[:50]}"
    )
    raise

# Save final model
print("\n[FINAL] Saving final model...")
final_path = f"{OUTPUT_DIR}/final"
trainer.save_model(final_path)
tokenizer.save_pretrained(final_path)

# Upload to HF
print("[FINAL] Uploading to Hugging Face...")
api.upload_folder(
    folder_path=final_path,
    repo_id="ZaoKing/khudi-ai-v2",
    commit_message="Khudi AI v2 final - trained on 42K Pakistani Q&A"
)
print(f"  ✅ Uploaded to: https://huggingface.co/ZaoKing/khudi-ai-v2")

# Cleanup
hc.stop()
print("\n🎉 V2 training complete!")
print(f"  Model: https://huggingface.co/ZaoKing/khudi-ai-v2")