"""
Evaluate the fine-tuned skill-classification model on Modal.

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

Run evaluation:
    modal run modal_apps/evaluate_modal.py

Reads eval_prompts.json locally, uploads it to the dataset volume, loads the
LoRA adapter from the model volume, and reports per-prompt PASS/FAIL plus accuracy.
"""

from __future__ import annotations

import json
import pathlib
import re

import modal

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

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

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
PROJECT_ROOT = pathlib.Path(__file__).resolve().parent.parent
LOCAL_EVAL_PROMPTS = PROJECT_ROOT / "data" / "eval_prompts.json"
REMOTE_EVAL_PROMPTS = "/data/eval_prompts.json"
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 = 30 * 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 (same as training)
# ---------------------------------------------------------------------------

eval_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 eval_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


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------


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

    eval_path = pathlib.Path(REMOTE_EVAL_PROMPTS)
    if not eval_path.exists():
        raise FileNotFoundError(
            f"Eval prompts not found at {eval_path}. "
            "Run `modal run modal_apps/evaluate_modal.py` from the project directory."
        )

    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."
        )

    with eval_path.open(encoding="utf-8") as handle:
        eval_prompts = json.load(handle)

    print(f"Loading base 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"Loading LoRA adapter from {ADAPTER_DIR}")
    model = PeftModel.from_pretrained(model, str(ADAPTER_DIR))

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

    FastLanguageModel.for_inference(model)

    passed = 0
    failed = 0
    total = len(eval_prompts)

    print(f"Running evaluation on {total} prompts...\n")

    for case in eval_prompts:
        prompt = case["prompt"]
        expected = case["expected"]

        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()
        predicted = extract_skill(raw_output)
        is_pass = predicted == expected

        if is_pass:
            passed += 1
        else:
            failed += 1

        print(f"Prompt: {prompt}")
        print(f"Expected: {expected}")
        print(f"Predicted: {predicted if predicted is not None else raw_output}")
        print(f"{'PASS' if is_pass else 'FAIL'}")
        print()

    accuracy = passed / total if total else 0.0

    print("--- Summary ---")
    print(f"Total: {total}")
    print(f"Passed: {passed}")
    print(f"Failed: {failed}")
    print(f"Accuracy: {accuracy:.1%}")


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


@app.local_entrypoint()
def main() -> None:
    eval_path = pathlib.Path(LOCAL_EVAL_PROMPTS)
    if not eval_path.exists():
        raise FileNotFoundError(
            f"Local eval prompts not found: {eval_path.resolve()}"
        )

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

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

    print("Launching evaluation on Modal GPU...")
    evaluate.remote()