"""
modal_app.py — Modal orchestration for LFED fine-tuning pipeline.

Deploy:  modal deploy modal_train/modal_app.py
Run:     modal run modal_train/modal_app.py

Pipeline:
  1. Generate synthetic NL→SQL pairs (CPU)
  2. QLoRA fine-tune on A10G (~3-6 hrs)
  3. Merge LoRA → GGUF Q4_K_M
  4. Push to Hugging Face Hub

Requires Modal secret 'huggingface' with HF_TOKEN.
"""

from pathlib import Path
import modal

# ── Modal app ──────────────────────────────────────────────────────────

app = modal.App("kasualdad-lfed-train")  # force redeploy v3

# ── Image with all dependencies ────────────────────────────────────────

train_image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("build-essential", "cmake", "git")
    .pip_install(
        "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git",
        "torch>=2.4.0",
        "transformers>=4.46.0",
        "trl>=0.12.0",
        "datasets>=3.0.0",
        "bitsandbytes>=0.44.0",
        "accelerate>=1.0.0",
        "peft>=0.13.0",
        "huggingface_hub>=0.26.0",
        "llama-cpp-python>=0.3.0",
        "gguf>=0.10.0",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
    .add_local_dir(
        Path(__file__).parent,
        remote_path="/root",
    )
)

# ── Volume for persistent storage (training data + checkpoints) ────────

volume = modal.Volume.from_name(
    "lfed-training-data", create_if_missing=True
)
MOUNT_DIR = Path("/data")


# ── Step 1: Generate synthetic data (CPU) ──────────────────────────────

@app.function(
    image=train_image,
    volumes={MOUNT_DIR: volume},
    timeout=10 * 60,  # 10 min
)
def generate_synthetic():
    """Generate NL→SQL training pairs."""
    import sys
    sys.path.insert(0, str(MOUNT_DIR))

    # Copy generate_synthetic.py to volume and run
    local_script = Path(__file__).parent / "generate_synthetic.py"
    dest = MOUNT_DIR / "generate_synthetic.py"
    dest.write_text(local_script.read_text())

    # Run generation
    from generate_synthetic import main
    main()

    # Copy output back
    import shutil
    train_jsonl = Path(__file__).parent / "train.jsonl"
    if train_jsonl.exists():
        shutil.move(str(train_jsonl), str(MOUNT_DIR / "train.jsonl"))

    print(f"✅ Synthetic data ready at {MOUNT_DIR / 'train.jsonl'}")
    volume.commit()


# ── Step 2: QLoRA fine-tuning (A10G) ───────────────────────────────────

@app.function(
    image=train_image,
    gpu="A10G",
    volumes={MOUNT_DIR: volume},
    timeout=6 * 3600,  # 6 hours
    secrets=[modal.Secret.from_name("huggingface")],
)
def train_model():
    """Fine-tune Qwen2.5-Coder-7B with QLoRA on A10G."""
    import sys
    sys.path.insert(0, str(MOUNT_DIR))

    # Copy train.py to volume and run
    train_script = Path(__file__).parent / "train.py"
    dest = MOUNT_DIR / "train.py"
    dest.write_text(train_script.read_text())

    # Ensure train.jsonl is available
    train_jsonl = MOUNT_DIR / "train.jsonl"
    if not train_jsonl.exists():
        # Try to copy from local
        local_jsonl = Path(__file__).parent / "train.jsonl"
        if local_jsonl.exists():
            train_jsonl.write_text(local_jsonl.read_text())
        else:
            raise FileNotFoundError(
                "train.jsonl not found. Run generate_synthetic first."
            )

    from train import main
    main()

    # Free GPU memory for the export step
    import gc
    import torch
    del sys.modules['train']
    gc.collect()
    torch.cuda.empty_cache()

    volume.commit()
    print("✅ Training complete")


# ── Step 3: Export GGUF + push to HF (A10G for merge speed) ────────────

@app.function(
    image=train_image,
    gpu="A10G",
    volumes={MOUNT_DIR: volume},
    timeout=2 * 3600,  # 2 hours
    secrets=[modal.Secret.from_name("huggingface")],
)
def export_and_push():
    """Merge LoRA, export GGUF, push to HF Hub."""
    import sys
    import importlib

    # Import directly from mount path (not volume copy)
    mount_path = str(Path(__file__).parent)
    if mount_path not in sys.path:
        sys.path.insert(0, mount_path)

    # Clear any cached module
    for key in list(sys.modules.keys()):
        if "export_gguf" in key:
            del sys.modules[key]

    import export_gguf
    importlib.reload(export_gguf)
    export_gguf.main()


# ── Full pipeline (run all steps sequentially) ─────────────────────────

@app.function(
    image=train_image,
    gpu="A10G",
    volumes={MOUNT_DIR: volume},
    timeout=8 * 3600,  # 8 hours total
    secrets=[modal.Secret.from_name("huggingface")],
)
def run_full_pipeline():
    """Run the complete pipeline: synthetic → train → export → push."""
    print("=" * 60)
    print("LFED Full Fine-Tuning Pipeline")
    print("=" * 60)

    print("\n📝 Step 1/3: Generating synthetic training data...")
    generate_synthetic.local()

    print("\n🏋️ Step 2/3: QLoRA fine-tuning on A10G...")
    train_model.local()

    print("\n📦 Step 3/3: Merging → GGUF → HF Hub...")
    export_and_push.remote()  # remote() = fresh container, clean GPU

    print("\n🎉 Full pipeline complete!")
    print(f"   Model repo: check your HF Hub for 'lfed-qwen2.5-coder-7b-sql-gguf'")


# ── Entry points ───────────────────────────────────────────────────────

@app.local_entrypoint()
def main():
    """
    Local entry point — kick off the full pipeline on Modal.

    Usage:
        modal run modal_train/modal_app.py
    """
    print("Deploying LFED fine-tuning pipeline to Modal...")
    run_full_pipeline.remote()