"""
Run skill-classification inference on Modal GPU.

Prerequisites:
    pip install modal
    modal setup
    modal run modal_apps/train_modal.py   # train and save LoRA adapter first

Run inference:
    modal run modal_apps/infer_modal.py --prompt "play my workout playlist"
"""

from __future__ import annotations

import json
import pathlib
import re

import modal

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

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

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
MODEL_DIR = pathlib.Path("/model")
ADAPTER_DIR = MODEL_DIR / "adapter"
MAX_SEQ_LENGTH = 2048

SYSTEM_PROMPT = (
    "You classify Android automation requests into exactly one skill. "
    'Reply with JSON only: {"skill": "<skill_name>"}. '
    "Use the app or action named in the request (contacts, Gmail, Slack, YouTube, etc.) "
    "to pick the correct skill."
)


def build_classifier_messages(user_content: str) -> list[dict[str, str]]:
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]

GPU_TYPE = "A10G"
TIMEOUT_SECONDS = 10 * 60

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

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 (same as training / evaluation)
# ---------------------------------------------------------------------------

infer_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 infer_image.imports():
    import unsloth  # noqa: F401 — must import before trl/transformers/peft

    import torch
    from peft import PeftModel
    from unsloth import FastLanguageModel
    from unsloth.chat_templates import get_chat_template


def extract_skill(text: str) -> str | None:
    """Extract the skill name from model output JSON."""
    text = text.strip()
    if not text:
        return None

    match = re.search(r'\{[^{}]*"skill"\s*:\s*"([^"]+)"[^{}]*\}', text)
    if match:
        return match.group(1)

    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end <= start:
        return None

    try:
        payload = json.loads(text[start : end + 1])
    except json.JSONDecodeError:
        return None

    skill = payload.get("skill")
    return skill if isinstance(skill, str) else None


def format_skill_json(skill: str) -> str:
    return json.dumps({"skill": skill}, separators=(",", ":"))


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


@app.function(
    image=infer_image,
    gpu=GPU_TYPE,
    timeout=TIMEOUT_SECONDS,
    volumes={
        "/model": model_volume,
        "/model_cache": model_cache_volume,
    },
)
def infer(prompt: str) -> str:
    model_volume.reload()

    if not (ADAPTER_DIR / "adapter_config.json").exists():
        raise FileNotFoundError(
            f"LoRA adapter not found at {ADAPTER_DIR}. "
            "Run `modal run modal_apps/train_modal.py` first."
        )

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=MAX_SEQ_LENGTH,
        dtype=None,
        load_in_4bit=True,
    )
    model = PeftModel.from_pretrained(model, str(ADAPTER_DIR))
    tokenizer = get_chat_template(
        tokenizer,
        chat_template="qwen-2.5",
    )
    FastLanguageModel.for_inference(model)

    messages = build_classifier_messages(prompt)
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.inference_mode():
        outputs = model.generate(
            input_ids=inputs,
            max_new_tokens=64,
            use_cache=True,
            do_sample=False,
        )

    generated = outputs[0][inputs.shape[1] :]
    raw_output = tokenizer.decode(generated, skip_special_tokens=True).strip()

    skill = extract_skill(raw_output)
    if skill is None:
        raise ValueError(f"Model did not return a skill for prompt: {prompt!r}")

    return format_skill_json(skill)


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


@app.local_entrypoint()
def main(prompt: str) -> None:
    print(infer.remote(prompt))