"""
Fine-tune Qwen/Qwen2.5-3B-Instruct on Modal with Unsloth (4-bit QLoRA + SFT).

Prerequisites (run once on your Mac):
    pip install modal
    modal setup

Run training:
    modal run modal_apps/train_modal.py

This uploads train.jsonl to a Modal Volume, trains on a GPU, and saves:
    /model/adapter  — LoRA adapter weights
    /model/merged   — merged 16-bit model

Download results after training:
    modal volume get android-dataset-model adapter ./trained_model/adapter
    modal volume get android-dataset-model merged ./trained_model/merged
"""

from __future__ import annotations

import pathlib

import modal

app = modal.App("android-skill-finetune")

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
PROJECT_ROOT = pathlib.Path(__file__).resolve().parent.parent
LOCAL_DATASET = PROJECT_ROOT / "data" / "train.jsonl"
REMOTE_DATASET = "/data/train.jsonl"
MODEL_DIR = pathlib.Path("/model")
ADAPTER_DIR = MODEL_DIR / "adapter"
MERGED_DIR = MODEL_DIR / "merged"
CHECKPOINT_DIR = MODEL_DIR / "checkpoints"
MAX_SEQ_LENGTH = 2048

LORA_R = 32
LORA_ALPHA = 32
NUM_EPOCHS = 5
BATCH_SIZE = 8

GPU_TYPE = "A10G"
TIMEOUT_SECONDS = 2 * 60 * 60

# ---------------------------------------------------------------------------
# Volumes
# ---------------------------------------------------------------------------

dataset_volume = modal.Volume.from_name(
    "android-dataset-data",
    create_if_missing=True,
)
model_volume = modal.Volume.from_name(
    "android-dataset-model",
    create_if_missing=True,
)
model_cache_volume = modal.Volume.from_name(
    "android-dataset-hf-cache",
    create_if_missing=True,
)

# ---------------------------------------------------------------------------
# Container image
# ---------------------------------------------------------------------------

train_image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install_from_requirements(
        str(pathlib.Path(__file__).parent / "requirements-modal.txt")
    )
    .env(
        {
            "HF_HOME": "/model_cache",
            "HF_HUB_ENABLE_HF_TRANSFER": "1",
        }
    )
)

with train_image.imports():
    import unsloth  # noqa: F401 — must import before trl/transformers/peft

    from datasets import load_dataset
    from trl import SFTConfig, SFTTrainer
    from unsloth import FastLanguageModel, is_bf16_supported
    from unsloth.chat_templates import get_chat_template


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------


@app.function(
    image=train_image,
    gpu=GPU_TYPE,
    timeout=TIMEOUT_SECONDS,
    volumes={
        "/data": dataset_volume,
        "/model": model_volume,
        "/model_cache": model_cache_volume,
    },
)
def train() -> None:
    dataset_volume.reload()

    data_path = pathlib.Path(REMOTE_DATASET)
    if not data_path.exists():
        raise FileNotFoundError(
            f"Dataset not found at {data_path}. "
            "Run `modal run modal_apps/train_modal.py` from the project directory."
        )

    print(f"Loading model: {MODEL_NAME}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=MAX_SEQ_LENGTH,
        dtype=None,
        load_in_4bit=True,
    )

    print(f"Applying QLoRA (rank={LORA_R})")
    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_R,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_alpha=LORA_ALPHA,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=3407,
        max_seq_length=MAX_SEQ_LENGTH,
    )

    tokenizer = get_chat_template(
        tokenizer,
        chat_template="qwen-2.5",
    )

    print(f"Loading dataset: {data_path}")
    dataset = load_dataset("json", data_files=str(data_path), split="train")
    print(f"Training examples: {len(dataset)}")

    def formatting_prompts_func(examples):
        texts = [
            tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=False,
            )
            for messages in examples["messages"]
        ]
        return {"text": texts}

    dataset = dataset.map(formatting_prompts_func, batched=True)

    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        args=SFTConfig(
            output_dir=str(CHECKPOINT_DIR),
            num_train_epochs=NUM_EPOCHS,
            per_device_train_batch_size=BATCH_SIZE,
            gradient_accumulation_steps=1,
            warmup_steps=10,
            learning_rate=2e-4,
            fp16=not is_bf16_supported(),
            bf16=is_bf16_supported(),
            logging_steps=10,
            optim="adamw_8bit",
            seed=3407,
            report_to="none",
            max_seq_length=MAX_SEQ_LENGTH,
            dataset_text_field="text",
            packing=False,
        ),
    )

    print("Starting training...")
    trainer.train()

    ADAPTER_DIR.mkdir(parents=True, exist_ok=True)
    MERGED_DIR.mkdir(parents=True, exist_ok=True)

    print(f"Saving LoRA adapter to {ADAPTER_DIR}")
    model.save_pretrained(str(ADAPTER_DIR))
    tokenizer.save_pretrained(str(ADAPTER_DIR))

    print(f"Saving merged 16-bit model to {MERGED_DIR}")
    model.save_pretrained_merged(
        str(MERGED_DIR),
        tokenizer,
        save_method="merged_16bit",
    )

    model_volume.commit()
    print("Training complete. Model saved to /model volume.")
    print(f"  Adapter: {ADAPTER_DIR}")
    print(f"  Merged:  {MERGED_DIR}")


# ---------------------------------------------------------------------------
# Local entrypoint
# ---------------------------------------------------------------------------


@app.local_entrypoint()
def main(dataset: str = "train_intent.jsonl") -> None:
    """Upload dataset and launch training.

    Args:
        dataset: Filename under data/ — use train_intent.jsonl for intent extraction.
    """
    dataset_path = PROJECT_ROOT / "data" / dataset
    if not dataset_path.exists():
        raise FileNotFoundError(
            f"Local dataset not found: {dataset_path.resolve()}"
        )

    remote_name = "train.jsonl"
    try:
        dataset_volume.remove_file(remote_name)
    except Exception:
        pass  # file may not exist yet on the volume

    print(f"Uploading {dataset_path} to dataset volume...")
    with dataset_volume.batch_upload() as batch:
        batch.put_file(str(dataset_path), remote_name)

    print("Launching training on Modal GPU...")
    train.remote()

    print()
    print("Done! Download your model with:")
    print("  modal volume get android-dataset-model adapter ./trained_model/adapter")
    print("  modal volume get android-dataset-model merged ./trained_model/merged")