#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch",
#     "peft>=0.18.0",
#     "transformers>=5.1.0",
#     "accelerate>=1.6.0",
#     "huggingface_hub",
# ]
# ///
"""
Merge LoRA Adapter and Push to Hub (HF Cloud)

This script is designed to run on Hugging Face Jobs infrastructure.
It loads a base model + LoRA adapter, merges them, and pushes the
standalone merged model to the Hub.

Usage:
    # Submit via Makefile
    make merge-cloud

    # Or submit directly
    python scripts/submit_cloud_training.py \
        --flavor a100x4 --timeout 2h \
        --script scripts/merge_hf_cloud.py
"""

import os

import torch
from huggingface_hub import login
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# ═══════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════

BASE_MODEL = "Qwen/Qwen3-Coder-Next"
ADAPTER_MODEL = "yangluo/univer-api-qwen3-next-coder"
MERGED_MODEL_ID = "yangluo/univer-api-qwen3-next-coder-merged"

# ═══════════════════════════════════════════════════════════════
# MERGE SCRIPT
# ═══════════════════════════════════════════════════════════════


def main():
    # Login to Hugging Face Hub
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        print("Logging in to Hugging Face Hub...")
        login(token=hf_token, add_to_git_credential=False)
        print("  Login successful!")
    else:
        print("ERROR: HF_TOKEN not found in environment!")
        print("Cannot push merged model without authentication.")
        return 1

    print("=" * 60)
    print("    MERGE LoRA ADAPTER & PUSH TO HUB")
    print("=" * 60)
    print(f"Base model:    {BASE_MODEL}")
    print(f"LoRA adapter:  {ADAPTER_MODEL}")
    print(f"Merged output: {MERGED_MODEL_ID}")
    print(f"Dtype:         bfloat16 (full precision, no quantization)")
    print("=" * 60)

    # ─────────────────────────────────────────────────────
    # Load Tokenizer
    # ─────────────────────────────────────────────────────
    print("\n[1/5] Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    print(f"  Tokenizer loaded: {tokenizer.__class__.__name__}")

    # ─────────────────────────────────────────────────────
    # Load Base Model (bf16, no quantization)
    # ─────────────────────────────────────────────────────
    print(f"\n[2/5] Loading base model: {BASE_MODEL}")
    print("  Loading in bfloat16 with device_map='auto'...")

    n_gpus = torch.cuda.device_count()
    print(f"  Detected {n_gpus} GPU(s)")

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )

    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        print(f"  Base model loaded, GPU memory: {allocated:.1f} GB")

    # ─────────────────────────────────────────────────────
    # Load LoRA Adapter
    # ─────────────────────────────────────────────────────
    print(f"\n[3/5] Loading LoRA adapter: {ADAPTER_MODEL}")
    model = PeftModel.from_pretrained(model, ADAPTER_MODEL)
    print("  Adapter loaded successfully")

    # ─────────────────────────────────────────────────────
    # Merge and Unload
    # ─────────────────────────────────────────────────────
    print("\n[4/5] Merging adapter into base model...")
    model = model.merge_and_unload()
    print("  Merge complete!")

    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        print(f"  GPU memory after merge: {allocated:.1f} GB")

    # ─────────────────────────────────────────────────────
    # Push to Hub
    # ─────────────────────────────────────────────────────
    print(f"\n[5/5] Pushing merged model to Hub: {MERGED_MODEL_ID}")
    print("  Pushing model (this may take a while for large models)...")
    model.push_to_hub(MERGED_MODEL_ID)
    print("  Model pushed!")

    print("  Pushing tokenizer...")
    tokenizer.push_to_hub(MERGED_MODEL_ID)
    print("  Tokenizer pushed!")

    print("\n" + "=" * 60)
    print("         MERGE COMPLETE!")
    print("=" * 60)
    print(f"Merged model: https://huggingface.co/{MERGED_MODEL_ID}")


if __name__ == "__main__":
    exit(main() or 0)