#!/usr/bin/env python3
"""Run the merged Gemma 4 Mobile Actions model with Transformers."""

from __future__ import annotations

import argparse
import json
import re
from typing import Any

import torch
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, pipeline


TOOL_CALL_RE = re.compile(
    r"<\|tool_call\>call:(\w+)\{(.*?)\}<(?:tool_call|tool)\|>",
    re.DOTALL,
)
STRING_ARG_RE = re.compile(r"(\w+):<\|\"\|>(.*?)<\|\"\|>", re.DOTALL)
PLAIN_ARG_RE = re.compile(r"(\w+):(None|True|False|-?\d+(?:\.\d+)?)")
QUOTED_VALUE_RE = re.compile(r":<\|\"\|>.*?<\|\"\|>", re.DOTALL)


TOOLS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "turn_on_flashlight",
            "description": "Turns on the device flashlight.",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "turn_off_flashlight",
            "description": "Turns off the device flashlight.",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "open_wifi_settings",
            "description": "Opens the device Wi-Fi settings screen.",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "show_map",
            "description": "Shows a location on the map.",
            "parameters": {
                "type": "object",
                "properties": {"query": {"type": "string"}},
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "send_email",
            "description": "Composes an email.",
            "parameters": {
                "type": "object",
                "properties": {
                    "recipient": {"type": "string"},
                    "subject": {"type": "string"},
                    "body": {"type": "string"},
                },
                "required": ["recipient", "body"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "create_calendar_event",
            "description": "Creates a calendar event.",
            "parameters": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "start_datetime": {"type": "string"},
                    "end_datetime": {"type": "string"},
                },
                "required": ["title", "start_datetime"],
            },
        },
    },
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-id", default="ClarkBear/gemma4-e2b-mobile-actions-200")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--system", default="You are a mobile assistant that calls tools.")
    parser.add_argument("--max-new-tokens", type=int, default=160)
    parser.add_argument("--dtype", choices=["auto", "bfloat16", "float16", "float32"], default="auto")
    return parser.parse_args()


def load_processor(model_id: str):
    try:
        return AutoProcessor.from_pretrained(model_id)
    except Exception:
        return AutoTokenizer.from_pretrained(model_id)


def tokenizer_from_processor(processor):
    return getattr(processor, "tokenizer", processor)


def torch_dtype(name: str) -> torch.dtype:
    if name == "float16":
        return torch.float16
    if name == "float32":
        return torch.float32
    return torch.bfloat16


def device_map():
    if torch.cuda.is_available():
        return "auto"
    if torch.backends.mps.is_available() and torch.backends.mps.is_built():
        return {"": "mps"}
    return {"": "cpu"}


def apply_template(processor, system: str, user_prompt: str) -> str:
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user_prompt},
    ]
    attempts = [
        {"tools": TOOLS, "add_generation_prompt": True, "enable_thinking": False},
        {"tools": TOOLS, "add_generation_prompt": True},
        {"add_generation_prompt": True, "enable_thinking": False},
        {"add_generation_prompt": True},
    ]
    last_error: Exception | None = None
    for kwargs in attempts:
        try:
            return processor.apply_chat_template(messages, tokenize=False, **kwargs)
        except TypeError as exc:
            last_error = exc
    raise RuntimeError(f"Could not apply chat template: {last_error}")


def parse_scalar(value: str) -> Any:
    if value == "None":
        return None
    if value == "True":
        return True
    if value == "False":
        return False
    if "." in value:
        return float(value)
    return int(value)


def parse_tool_call(text: str) -> dict[str, Any] | None:
    match = TOOL_CALL_RE.search(text)
    if not match:
        return None
    name, body = match.group(1), match.group(2)
    args: dict[str, Any] = {}
    for key, value in STRING_ARG_RE.findall(body):
        args[key] = value
    body_without_strings = QUOTED_VALUE_RE.sub("", body)
    for key, value in PLAIN_ARG_RE.findall(body_without_strings):
        args.setdefault(key, parse_scalar(value))
    return {"name": name, "args": args, "raw": match.group(0)}


def main() -> None:
    args = parse_args()
    processor = load_processor(args.model_id)
    tokenizer = tokenizer_from_processor(processor)
    prompt = apply_template(processor, args.system, args.prompt)

    model = AutoModelForCausalLM.from_pretrained(
        args.model_id,
        dtype=torch_dtype(args.dtype),
        device_map=device_map(),
    )
    generator = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        clean_up_tokenization_spaces=False,
    )
    output = generator(
        prompt,
        max_new_tokens=args.max_new_tokens,
        do_sample=False,
    )[0]["generated_text"]
    generated = output[len(prompt) :]

    print("=== Generated ===")
    print(generated.strip())
    print("\n=== Parsed Tool Call ===")
    print(json.dumps(parse_tool_call(generated), indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()