"""
Evaluate the fine-tuned intent-extraction model on Modal.

Prerequisites:
    pip install modal
    modal setup
    python scripts/generate_intent_dataset.py
    modal run modal_apps/train_modal.py   # train on data/train_intent.jsonl

Run evaluation:
    modal run modal_apps/evaluate_intent_modal.py

Reads eval_intent_prompts.json locally, uploads it to the dataset volume, loads the
LoRA adapter from the model volume, and reports per-prompt PASS/FAIL for skill +
parameters plus overall accuracy.
"""

from __future__ import annotations

import json
import pathlib

import modal

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

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

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
PROJECT_ROOT = pathlib.Path(__file__).resolve().parent.parent
LOCAL_EVAL_PROMPTS = PROJECT_ROOT / "data" / "eval_intent_prompts.json"
REMOTE_EVAL_PROMPTS = "/data/eval_intent_prompts.json"
MODEL_DIR = pathlib.Path("/model")
ADAPTER_DIR = MODEL_DIR / "adapter"
MAX_SEQ_LENGTH = 2048
MAX_NEW_TOKENS = 128

INTENT_SYSTEM_PROMPT = (
    "You extract structured Android automation intents from natural language. "
    'Reply with JSON only: {"skill": "<skill_name>", "parameters": {<extracted_fields>}}. '
    "Pick exactly one skill. Extract all relevant parameters mentioned in the request "
    "(contact names, messages, times, destinations, channel names, search queries, etc.). "
    "Use an empty object for parameters when the skill needs none. "
    "Use the app or action named in the request (contacts, Gmail, Slack, YouTube, etc.) "
    "to pick the correct skill."
)


def build_intent_messages(user_content: str) -> list[dict[str, str]]:
    return [
        {"role": "system", "content": INTENT_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
# ---------------------------------------------------------------------------

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 _parse_json_payload(text: str) -> dict | None:
    text = text.strip()
    if not text:
        return None

    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

    return payload if isinstance(payload, dict) else None


def extract_intent(text: str) -> dict | None:
    payload = _parse_json_payload(text)
    if not payload:
        return None

    skill = payload.get("skill")
    if not isinstance(skill, str) or not skill:
        return None

    parameters = payload.get("parameters", {})
    if parameters is None:
        parameters = {}
    if not isinstance(parameters, dict):
        return None

    return {"skill": skill, "parameters": parameters}


def normalize_param(value: str) -> str:
    return " ".join(value.lower().strip().split())


def parameters_match(predicted: dict, expected: dict) -> bool:
    for key, expected_value in expected.items():
        predicted_value = predicted.get(key)
        if predicted_value is None:
            return False
        if normalize_param(str(predicted_value)) != normalize_param(str(expected_value)):
            return False
    return True


def intent_matches(predicted: dict | None, expected: dict) -> tuple[bool, bool]:
    """Return (skill_match, full_match)."""
    if not predicted:
        return False, False

    expected_skill = expected["skill"]
    expected_params = expected.get("parameters", {})

    skill_match = predicted.get("skill") == expected_skill
    if not skill_match:
        return False, False

    if not expected_params:
        return True, True

    params_match = parameters_match(predicted.get("parameters", {}), expected_params)
    return True, params_match


# ---------------------------------------------------------------------------
# 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_intent_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)

    skill_passed = 0
    full_passed = 0
    total = len(eval_prompts)

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

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

        messages = build_intent_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=MAX_NEW_TOKENS,
                use_cache=True,
                do_sample=False,
            )

        generated = outputs[0][inputs.shape[1] :]
        raw_output = tokenizer.decode(generated, skip_special_tokens=True).strip()
        predicted = extract_intent(raw_output)
        skill_ok, full_ok = intent_matches(predicted, expected)

        if skill_ok:
            skill_passed += 1
        if full_ok:
            full_passed += 1

        print(f"Prompt: {prompt}")
        print(f"Expected: {json.dumps(expected, separators=(',', ':'))}")
        print(f"Predicted: {json.dumps(predicted, separators=(',', ':')) if predicted else raw_output}")
        print(f"Skill: {'PASS' if skill_ok else 'FAIL'} | Full: {'PASS' if full_ok else 'FAIL'}")
        print()

    skill_accuracy = skill_passed / total if total else 0.0
    full_accuracy = full_passed / total if total else 0.0

    print("--- Summary ---")
    print(f"Total: {total}")
    print(f"Skill accuracy: {skill_passed}/{total} ({skill_accuracy:.1%})")
    print(f"Full intent accuracy: {full_passed}/{total} ({full_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()}. "
            "Run `python scripts/generate_intent_dataset.py` first."
        )

    remote_name = "eval_intent_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 intent evaluation on Modal GPU...")
    evaluate.remote()