#!/bin/bash
# Khudi AI v2 — SMOKE TEST (30 min)
# Tests the full pipeline before launching full training

set -e
exec > >(tee -a /workspace/smoke.log) 2>&1

echo "============================================"
echo "🧪 KHUDI v2 SMOKE TEST (30 min)"
echo "============================================"
echo "Start: $(date)"

# Install packages
echo "[1/5] Installing packages..."
pip install -q --no-cache-dir \
    "transformers==4.45.0" \
    "datasets==2.20.0" \
    "peft==0.11.0" \
    "trl==0.10.0" \
    "accelerate==0.34.0" \
    "bitsandbytes==0.43.3" \
    "huggingface_hub==0.25.0" \
    "sentencepiece" "protobuf" 2>&1 | tail -3

# Setup dirs
mkdir -p /workspace
cd /workspace
mkdir -p khudi-v2-output

# Download SMOKE TEST dataset (just 100 samples)
echo "[2/5] Downloading smoke test data (100 samples)..."
wget -q --header="Authorization: Bearer HF_TOKEN_PLACEHOLDER" \
     -O smoke_data.jsonl \
     "https://huggingface.co/datasets/ZaoKing/khudi-ai-dataset/resolve/main/v2_chunk_0.jsonl"

# Take only first 100 lines
head -100 smoke_data.jsonl > smoke_data_small.jsonl
mv smoke_data_small.jsonl smoke_data.jsonl
echo "✅ $(wc -l < smoke_data.jsonl) samples ready"

# Run smoke test (5 steps)
echo "[3/5] Running smoke test (5 steps)..."
python -c "
import os
os.environ['HF_TOKEN'] = 'HF_TOKEN_PLACEHOLDER'
os.environ['HUGGINGFACE_HUB_TOKEN'] = 'HF_TOKEN_PLACEHOLDER'
os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

print('Loading model...')
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type='nf4',
                        bnb_4bit_compute_dtype=torch.bfloat16,
                        bnb_4bit_use_double_quant=True)
tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-7B-Instruct', trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-7B-Instruct',
    quantization_config=bnb, device_map='auto', trust_remote_code=True)
model.config.use_cache = False
model = prepare_model_for_kbit_training(model)
lora = 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)
model.print_trainable_parameters()

print('Loading data...')
ds = load_dataset('json', data_files='smoke_data.jsonl', split='train')
def fmt(ex):
    return {'text': tokenizer.apply_chat_template(ex['messages'], tokenize=False)}
ds = ds.map(fmt, remove_columns=ds.column_names)
print(f'Data: {len(ds)} samples')

print('Training (5 steps)...')
args = SFTConfig(
    output_dir='/workspace/khudi-v2-output/smoke',
    max_steps=5, per_device_train_batch_size=2,
    gradient_accumulation_steps=2,
    max_length=512, learning_rate=2e-4,
    logging_steps=1, save_strategy='no',
    bf16=True, report_to='none',
    optim='paged_adamw_8bit',
)
trainer = SFTTrainer(model=model, args=args, train_dataset=ds, processing_class=tokenizer)
trainer.train()
print('✅ SMOKE TEST PASSED!')
" 2>&1

# Save smoke test result to HF
echo "[4/5] Saving smoke test results..."
python -c "
import os
from huggingface_hub import HfApi
api = HfApi()
with open('/workspace/smoke.log', 'r') as f:
    log = f.read()
if 'SMOKE TEST PASSED' in log:
    api.upload_file(
        path_or_fileobj='/workspace/smoke.log',
        path_in_repo='smoke_test_passed.log',
        repo_id='ZaoKing/khudi-ai-v2',
        repo_type='model',
        commit_message='Smoke test PASSED - safe to launch full training'
    )
    print('✅ Smoke test log uploaded')
else:
    api.upload_file(
        path_or_fileobj='/workspace/smoke.log',
        path_in_repo='smoke_test_FAILED.log',
        repo_id='ZaoKing/khudi-ai-v2',
        repo_type='model',
        commit_message='Smoke test FAILED - check logs'
    )
    print('❌ Smoke test failed log uploaded')
" 2>&1

echo "[5/5] Smoke test complete: $(date)"
echo "============================================"