"""
╔══════════════════════════════════════════════════════════════╗
║     MamunAI — Hugging Face Spaces (ZeroGPU) Training       ║
║     Owner : Al Mamun Khan                                   ║
║                                                             ║
║  HOW TO USE:                                                ║
║  1. Create a new HF Space (Gradio, ZeroGPU hardware)       ║
║  2. Upload: hf_space_train.py, dataset.jsonl, requirements ║
║  3. Click "Run Training" button in the UI                   ║
║  4. Download GGUF from Files tab when done                  ║
╚══════════════════════════════════════════════════════════════╝
"""

import gradio as gr
import json, os, shutil, subprocess, threading, time
from pathlib import Path

# ── Config ────────────────────────────────────────────────────────────────────
BASE_MODEL   = "HuggingFaceTB/SmolLM-360M-Instruct"
ADAPTER_DIR  = "./lora_adapter"
MERGED_DIR   = "./merged_model"
GGUF_OUT     = "./MamunAI-finetuned.gguf"
DATASET_FILE = "./dataset.jsonl"

LORA_RANK    = 16
LORA_ALPHA   = 32
LORA_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]
BATCH_SIZE   = 16    # A100 → can go higher; T4 free tier → 16 safe
EPOCHS       = 3
LR           = 2e-4
MAX_SEQ_LEN  = 256

SYSTEM = (
    "You are an Android phone control assistant. You do not possess any internal data. "
    "Your sole function is to convert user input into JSON commands. "
    "You understand English. You have access to tools:\n"
    "1. android_control\n2. web_search\n3. direct_answer\n4. clarify\n\n"
    "Your output must be strictly in JSON format.\n\n"
    'Example: User: "Dim the lights" -> Output: {"tool": "android_control", "action": "set_brightness", "params": {"value": 30}}\n'
    'Example: User: "Today\'s weather?" -> Output: {"tool": "web_search", "query": "Dhaka weather 2026"}\n'
    'Example: User: "Hello" -> Output: {"tool": "direct_answer", "response": "I am doing well!"}\n\n'
    "Now process:"
)

# ── Training log stream ───────────────────────────────────────────────────────
_log_lines = []
_training_done = False
_gguf_ready = False

def log(msg):
    _log_lines.append(msg)
    print(msg)

def get_logs():
    return "\n".join(_log_lines[-80:])   # last 80 lines

# ── Core training function ────────────────────────────────────────────────────
def run_training(dataset_file):
    global _training_done, _gguf_ready, _log_lines
    _log_lines = []
    _training_done = False
    _gguf_ready = False

    try:
        import torch
        from transformers import (AutoTokenizer, AutoModelForCausalLM,
                                  BitsAndBytesConfig, TrainingArguments)
        from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
        from trl import SFTTrainer
        from datasets import Dataset

        # ── GPU check ──────────────────────────────────────────────────────
        if not torch.cuda.is_available():
            log("❌ No GPU! Please use ZeroGPU Space or T4 runtime.")
            return
        log(f"✅ GPU: {torch.cuda.get_device_name(0)}")
        log(f"   VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB")

        # ── Load dataset ───────────────────────────────────────────────────
        ds_path = dataset_file if dataset_file else DATASET_FILE
        texts = []
        with open(ds_path) as f:
            for line in f:
                line = line.strip()
                if not line: continue
                rec = json.loads(line)
                text = (
                    f"<|im_start|>system\n{SYSTEM.strip()}<|im_end|>\n"
                    f"<|im_start|>user\nUser: \"{rec['instruction']}\"<|im_end|>\n"
                    f"<|im_start|>assistant\n{rec['output']}<|im_end|>"
                )
                texts.append(text)
        log(f"📂 Loaded {len(texts):,} examples")

        # ── Tokenizer ──────────────────────────────────────────────────────
        log(f"⬇️  Loading tokenizer...")
        from transformers import AutoTokenizer
        tok = AutoTokenizer.from_pretrained(BASE_MODEL)
        if tok.pad_token is None:
            tok.pad_token = tok.eos_token
        tok.padding_side = "right"

        # ── Model (4-bit QLoRA) ────────────────────────────────────────────
        log(f"⬇️  Loading model (4-bit QLoRA)...")
        bnb = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.float16,
        )
        model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL, quantization_config=bnb,
            device_map="auto", low_cpu_mem_usage=True,
        )
        model.config.use_cache = False

        # ── LoRA ───────────────────────────────────────────────────────────
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
        lora_cfg = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=LORA_RANK, lora_alpha=LORA_ALPHA,
            lora_dropout=0.05, target_modules=LORA_MODULES,
            bias="none", inference_mode=False,
        )
        model = get_peft_model(model, lora_cfg)
        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
        log(f"🔧 Trainable params: {trainable/1e6:.2f}M")

        # ── Train ──────────────────────────────────────────────────────────
        os.makedirs(ADAPTER_DIR, exist_ok=True)
        dataset = Dataset.from_dict({"text": texts})
        args = TrainingArguments(
            output_dir=ADAPTER_DIR,
            num_train_epochs=EPOCHS,
            per_device_train_batch_size=BATCH_SIZE,
            gradient_accumulation_steps=2,
            gradient_checkpointing=True,
            warmup_ratio=0.03,
            learning_rate=LR,
            lr_scheduler_type="cosine",
            fp16=True, bf16=False,
            optim="paged_adamw_8bit",
            logging_steps=10,
            save_steps=500,
            save_total_limit=1,
            report_to="none",
            group_by_length=True,
        )
        trainer = SFTTrainer(
            model=model, tokenizer=tok,
            train_dataset=dataset,
            dataset_text_field="text",
            max_seq_length=MAX_SEQ_LEN,
            packing=True, args=args,
        )
        log("🚀 Training started...")
        trainer.train()

        # ── Save adapter ───────────────────────────────────────────────────
        model.save_pretrained(ADAPTER_DIR)
        tok.save_pretrained(ADAPTER_DIR)
        log(f"💾 Adapter saved → {ADAPTER_DIR}/")

        # ── Merge ──────────────────────────────────────────────────────────
        log("🔀 Merging LoRA into base model...")
        from peft import PeftModel
        base = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL, torch_dtype=torch.float16, low_cpu_mem_usage=True
        )
        peft_m = PeftModel.from_pretrained(base, ADAPTER_DIR)
        merged = peft_m.merge_and_unload()
        os.makedirs(MERGED_DIR, exist_ok=True)
        merged.save_pretrained(MERGED_DIR)
        tok.save_pretrained(MERGED_DIR)
        log(f"✅ Merged model → {MERGED_DIR}/")

        # ── GGUF ───────────────────────────────────────────────────────────
        log("⚙️  Converting to GGUF Q4_K_M...")
        llama_dir = "/tmp/llama.cpp"
        if not Path(llama_dir).exists():
            subprocess.run(["git","clone","--depth","1",
                            "https://github.com/ggerganov/llama.cpp", llama_dir],
                           check=True, capture_output=True)
        f16 = "/tmp/mamunai-f16.gguf"
        subprocess.run(
            ["python3", f"{llama_dir}/convert_hf_to_gguf.py",
             MERGED_DIR, "--outfile", f16, "--outtype", "f16"],
            check=True
        )
        import ctypes, llama_cpp as lc
        params = lc.llama_model_quantize_default_params()
        params.ftype = 15   # Q4_K_M
        params.nthread = 4
        lc.llama_model_quantize(f16.encode(), GGUF_OUT.encode(), ctypes.byref(params))
        Path(f16).unlink(missing_ok=True)
        size = Path(GGUF_OUT).stat().st_size / 1e6
        log(f"✅ GGUF ready → {GGUF_OUT}  ({size:.0f} MB)")

        _gguf_ready = True
        _training_done = True
        log("🎉 ALL DONE! Download MamunAI-finetuned.gguf from Files tab.")

    except Exception as e:
        import traceback
        log(f"❌ Error: {e}")
        log(traceback.format_exc())
        _training_done = True


# ── Gradio UI ─────────────────────────────────────────────────────────────────
def start_training(dataset_upload):
    global _log_lines, _training_done
    _log_lines = ["⏳ Starting training..."]
    _training_done = False

    ds_path = dataset_upload if dataset_upload else DATASET_FILE
    if not Path(ds_path).exists():
        return "❌ dataset.jsonl not found. Please upload it.", None

    t = threading.Thread(target=run_training, args=(ds_path,), daemon=True)
    t.start()
    return "✅ Training started! Logs updating below...", None

def refresh_logs():
    done_msg = "\n\n✅ TRAINING COMPLETE — download GGUF from Files tab!" if _training_done else ""
    return get_logs() + done_msg

def download_gguf():
    if Path(GGUF_OUT).exists():
        return GGUF_OUT
    return None

with gr.Blocks(title="MamunAI Trainer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# 🤖 MamunAI Trainer — Hugging Face Space
**Owner:** Al Mamun Khan | Model: SmolLM-360M-Instruct + LoRA Q4_K_M

> Upload your `dataset.jsonl` and click **Start Training**.
> Training takes ~45–60 min on A100, ~60–90 min on T4.
""")

    with gr.Row():
        with gr.Column(scale=1):
            dataset_file = gr.File(
                label="📂 Upload dataset.jsonl",
                file_types=[".jsonl"],
            )
            train_btn  = gr.Button("🚀 Start Training", variant="primary", size="lg")
            status_box = gr.Textbox(label="Status", lines=2, interactive=False)
            refresh_btn = gr.Button("🔄 Refresh Logs", size="sm")
            download_btn = gr.DownloadButton(
                label="⬇️ Download MamunAI-finetuned.gguf",
                visible=True,
            )

        with gr.Column(scale=2):
            log_box = gr.Textbox(
                label="📋 Training Logs",
                lines=30,
                max_lines=30,
                interactive=False,
                autoscroll=True,
            )

    # ── Wiring ────────────────────────────────────────────────────────────────
    train_btn.click(
        fn=start_training,
        inputs=[dataset_file],
        outputs=[status_box, log_box],
    )
    refresh_btn.click(fn=refresh_logs, outputs=[log_box])
    download_btn.click(fn=download_gguf, outputs=[download_btn])

    gr.Markdown("""
---
### ⚙️ Training Config
| Setting | Value |
|---------|-------|
| Base Model | SmolLM-360M-Instruct |
| LoRA Rank | 16 |
| Target Modules | q, k, v, o projections |
| Batch Size | 16 |
| Epochs | 3 |
| Optimizer | paged_adamw_8bit |
| Precision | FP16 |
| Output | Q4_K_M GGUF (~260 MB) |
""")

if __name__ == "__main__":
    demo.launch()