"""
Khudi AI v2 — Production training script.
Designed to be uploaded to vast.ai instance and run.
All safety features built-in.
"""
import os
import sys
import time
import json
import argparse
import threading
import subprocess
from pathlib import Path

# Configuration
MODEL_NAME = os.environ.get("MODEL_NAME", "Qwen/Qwen3.5-9B")
DATA_PATH = os.environ.get("DATA_PATH", "/workspace/v2_data.jsonl")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/workspace/khudi-v2-output")
HF_TOKEN = os.environ.get("HF_TOKEN", "")
REPO_ID = "ZaoKing/khudi-ai-v2"

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

# ============================================================
# HEALTH CHECK
# ============================================================
class HealthMonitor:
    def __init__(self):
        self.last_active = time.time()
        self.stop_flag = False
        self.upload_count = 0

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

    def loop(self):
        while not self.stop_flag:
            time.sleep(60)
            util = self.get_gpu_util()
            if util > 5:
                self.last_active = time.time()
            idle_min = (time.time() - self.last_active) / 60
            # Log to file
            with open(f"{OUTPUT_DIR}/health.log", "a") as f:
                f.write(f"[{time.strftime('%H:%M:%S')}] GPU:{util}% idle:{idle_min:.1f}min\n")
            if idle_min > 15:
                print(f"\n[HEALTH] GPU idle {idle_min:.0f}min - emergency exit")
                # Save what we have
                self.emergency_save()
                os._exit(2)

    def emergency_save(self):
        """Try to save current state"""
        try:
            from huggingface_hub import HfApi
            api = HfApi()
            # Find latest checkpoint
            ckpt_dir = f"{OUTPUT_DIR}/checkpoints"
            if os.path.exists(ckpt_dir):
                checkpoints = [d for d in os.listdir(ckpt_dir) if d.startswith("checkpoint-")]
                if checkpoints:
                    checkpoints.sort(key=lambda x: int(x.split("-")[1]))
                    latest = checkpoints[-1]
                    api.upload_folder(
                        folder_path=f"{ckpt_dir}/{latest}",
                        repo_id=REPO_ID,
                        repo_type="model",
                        commit_message=f"EMERGENCY: {latest} - health check timeout"
                    )
                    print(f"✅ Emergency saved {latest}")
        except Exception as e:
            print(f"❌ Emergency save failed: {e}")

    def start(self):
        t = threading.Thread(target=self.loop, daemon=True)
        t.start()

    def stop(self):
        self.stop_flag = True


# ============================================================
# MAIN
# ============================================================
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--resume_from_checkpoint", action="store_true")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--grad_accum", type=int, default=16)
    parser.add_argument("--max_length", type=int, default=512)
    parser.add_argument("--lr", type=float, default=2e-4)
    args = parser.parse_args()

    print("=" * 70)
    print("🇵🇰 KHUDI AI v2 — TRAINING (Safe Mode)")
    print("=" * 70)
    print(f"Model: {MODEL_NAME}")
    print(f"Data: {DATA_PATH}")
    print(f"Output: {OUTPUT_DIR}")
    print(f"HF Repo: {REPO_ID}")
    print(f"Epochs: {args.epochs}, Batch: {args.batch_size}, Grad accum: {args.grad_accum}")
    print()

    # Imports after env setup
    import torch
    from datasets import load_dataset
    from transformers import (
        AutoModelForCausalLM, AutoTokenizer,
        BitsAndBytesConfig, TrainingArguments,
    )
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from trl import SFTTrainer, SFTConfig
    from huggingface_hub import HfApi, login

    # Login
    if HF_TOKEN:
        login(token=HF_TOKEN, add_to_git_credential=False)
        print("✅ Logged in to HF")

    api = HfApi()

    # ============================================================
    # LOAD DATASET
    # ============================================================
    print("[1/6] Loading dataset...")
    if not os.path.exists(DATA_PATH):
        print(f"❌ Dataset not found: {DATA_PATH}")
        sys.exit(1)

    dataset = load_dataset("json", data_files=DATA_PATH, split="train")
    print(f"  Loaded {len(dataset)} samples")

    # Filter for reasonable length
    def length_filter(example):
        total = sum(len(m.get("content", "")) for m in example.get("messages", []))
        return 100 < total < 4000
    dataset = dataset.filter(length_filter)
    print(f"  After length filter: {len(dataset)} samples")

    # ============================================================
    # LOAD TOKENIZER
    # ============================================================
    print("[2/6] Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # Format with chat template
    def format_chat(example):
        msgs = example.get("messages", [])
        text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
        return {"text": text}
    dataset = dataset.map(format_chat, remove_columns=dataset.column_names)
    print(f"  ✅ Formatted {len(dataset)} samples")

    # ============================================================
    # LOAD MODEL (4-bit)
    # ============================================================
    print("[3/6] Loading model (4-bit)...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        attn_implementation="sdpa",
    )
    model.config.use_cache = False
    model = prepare_model_for_kbit_training(model)

    # LoRA
    lora_config = LoraConfig(
        r=32, 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()

    # ============================================================
    # TRAINING CONFIG
    # ============================================================
    print("[4/6] Setting up trainer...")

    training_args = SFTConfig(
        output_dir=f"{OUTPUT_DIR}/checkpoints",
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        max_length=args.max_length,
        logging_steps=10,
        save_steps=50,            # ⭐ More frequent (was 100, now 50 for safety)
        save_total_limit=2,       # Keep 2 most recent
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        optim="paged_adamw_8bit",
        bf16=True,
        gradient_checkpointing=True,
        report_to="none",
        push_to_hub=True,         # ⭐ Auto-push to HF
        hub_model_id=REPO_ID,
        hub_token=HF_TOKEN,
        hub_strategy="checkpoint",
        save_safetensors=True,
        dataloader_num_workers=2, # ⭐ Speed up data loading
        remove_unused_columns=False,  # ⭐ Don't lose data
    )

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

    # ============================================================
    # START HEALTH MONITOR
    # ============================================================
    print("[5/6] Starting health monitor...")
    monitor = HealthMonitor()
    monitor.start()

    # ============================================================
    # TRAIN
    # ============================================================
    print("[6/6] Training...")
    print("=" * 70)

    try:
        if args.resume_from_checkpoint:
            trainer.train(resume_from_checkpoint=True)
        else:
            trainer.train()
        print("\n✅ Training completed!")
    except Exception as e:
        print(f"\n❌ Training error: {e}")
        monitor.emergency_save()
        raise

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

    # Upload final to HF
    print("[FINAL] Uploading to HF...")
    api.upload_folder(
        folder_path=final_path,
        repo_id=REPO_ID, repo_type="model",
        commit_message="Khudi AI v2 - trained on 42K Pakistani Q&A"
    )
    print(f"  ✅ Model: https://huggingface.co/{REPO_ID}")

    monitor.stop()
    print("\n🎉 V2 training complete!")


if __name__ == "__main__":
    main()