"""
Evaluate the Pocket Automator real-world benchmark on Modal.

Prerequisites:
    pip install modal
    modal setup
    python scripts/generate_pocket_benchmark.py
    modal run modal_apps/train_modal.py

Run evaluation:
    modal run modal_apps/evaluate_pocket_benchmark_modal.py
"""

from __future__ import annotations

import json
import pathlib

import modal

app = modal.App("pocket-automator-benchmark")

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
PROJECT_ROOT = pathlib.Path(__file__).resolve().parent.parent
LOCAL_BENCHMARK_PROMPTS = PROJECT_ROOT / "data" / "pocket_benchmark_prompts.json"
REMOTE_BENCHMARK_PROMPTS = "/data/pocket_benchmark_prompts.json"
REMOTE_RESULTS = "/data/pocket_benchmark_results.json"
REMOTE_REPORT = "/data/pocket_benchmark_report.txt"
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."
)

GPU_TYPE = "A10G"
TIMEOUT_SECONDS = 45 * 60

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,
)

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",
            "PYTHONPATH": "/root/src",
        }
    )
    .add_local_dir(str(PROJECT_ROOT / "src"), remote_path="/root/src", copy=True)
)


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


with eval_image.imports():
    import unsloth  # noqa: F401

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

    from pocket_benchmark import record_result, save_benchmark_outputs
    from skill_utils import extract_intent


@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()

    benchmark_path = pathlib.Path(REMOTE_BENCHMARK_PROMPTS)
    if not benchmark_path.exists():
        raise FileNotFoundError(
            f"Benchmark prompts not found at {benchmark_path}. "
            "Run `modal run modal_apps/evaluate_pocket_benchmark_modal.py` locally first."
        )

    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 benchmark_path.open(encoding="utf-8") as handle:
        cases = 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)

    results = []
    total = len(cases)
    print(f"Running Pocket Automator benchmark on {total} prompts...\n")

    for index, case in enumerate(cases, start=1):
        prompt = case["prompt"]
        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)
        result = record_result(case, raw_output, predicted)
        results.append(result)

        print(f"--- [{index}/{total}] {case.get('id', 'n/a')} ---")
        print(f"Prompt: {prompt}")
        print(f"Expected: {json.dumps(case['expected'], separators=(',', ':'))}")
        print(
            f"Predicted: {json.dumps(predicted, separators=(',', ':')) if predicted else raw_output}"
        )
        print(
            f"Skill: {'PASS' if result['skill_correct'] else 'FAIL'} | "
            f"Params: {'PASS' if result['parameter_correct'] else 'FAIL'} | "
            f"Exact JSON: {'PASS' if result['exact_json_match'] else 'FAIL'}"
        )
        print()

    metrics, report = save_benchmark_outputs(
        results,
        results_path=pathlib.Path(REMOTE_RESULTS),
        report_path=pathlib.Path(REMOTE_REPORT),
    )
    dataset_volume.commit()

    print("--- Pocket Automator Benchmark Summary ---")
    print(report)
    print(f"Saved results to {REMOTE_RESULTS}")
    print(f"Saved report to {REMOTE_REPORT}")


@app.local_entrypoint()
def main() -> None:
    benchmark_path = pathlib.Path(LOCAL_BENCHMARK_PROMPTS)
    if not benchmark_path.exists():
        raise FileNotFoundError(
            f"Local benchmark prompts not found: {benchmark_path.resolve()}. "
            "Run `python scripts/generate_pocket_benchmark.py` first."
        )

    remote_name = "pocket_benchmark_prompts.json"
    try:
        dataset_volume.remove_file(remote_name)
    except Exception:
        pass

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

    print("Launching Pocket Automator benchmark on Modal GPU...")
    evaluate.remote()

    print(
        "Benchmark complete. Results saved to Modal volume 'android-dataset-data':\n"
        "  - pocket_benchmark_results.json\n"
        "  - pocket_benchmark_report.txt"
    )