"""Versioned Android-only 105-slot adapter; frozen 133-slot models are unchanged.

The user-selected layout is action(1), two coordinate pairs(4), text(100).
Reuse the pinned GUI-Decisions coordinate codebook/targets/readout and the
existing exact text/EOS codec. The one-slot action code is a new, explicit
categorical mapping, not a claim that action names tokenize to one token.
"""

import hashlib
import json
from dataclasses import dataclass, field, replace

from trace2task.mobile133_codec import BUNDLE, Mobile133Codec
from trace2task.mobile133_protocol import MobileExecutionConfig

PROTOCOL = "gui-swift-mobile105-v1"
ACTIONS = (
    "tap",
    "long_press",
    "swipe",
    "scroll_up",
    "scroll_down",
    "scroll_left",
    "scroll_right",
    "type_text",
    "open_app",
    "back",
    "home",
    "enter",
    "menu",
    "wait",
    "done",
    "failed",
)
ACTION_LABELS = tuple("ABCDEFGHIJKLMNOP")
FIELD_WIDTHS = (
    ("action", 1),
    ("coordinate.x", 1),
    ("coordinate.y", 1),
    ("coordinate2.x", 1),
    ("coordinate2.y", 1),
    ("text", 100),
)
LITERALS = (
    '<action>{"action":"',
    '","coordinate":[',
    ",",
    '],"coordinate2":[',
    ",",
    '],"text":"',
    '"}</action>',
)
# Only coordinate/text value positions are copied from the pinned parameter codec.
PARAMETER_POSITIONS = {**{p: p - 1 for p in range(2, 6)}, **{p: p - 28 for p in range(33, 133)}}
SYSTEM = (
    "Execute the current Android instruction using the screenshot and executed-action history. "
    "Predict the 105 masked values in one forward: action, x, y, x2, y2, and 100 text slots. "
    "Action codes: " + ", ".join(f"{label}={a}" for label, a in zip(ACTION_LABELS, ACTIONS)) + ". "
    "tap/long_press use x,y; swipe uses both coordinate pairs; type_text/open_app use text. "
    "Other actions have no parameters. Coordinates and historical coordinates use the same "
    "256-bin codebook; the decoder maps them to 0-1000. Scroll direction means browsing direction. "
    "Copy text exactly, at most 99 native tokens followed by EOS. Timing is an executor setting. "
    "Predict only applicable fields in the supplied <action> table; do not output explanations."
)


def from_source_action(action):
    """Lossless semantic projection of an admitted Android action, also for history."""
    a = action.get("action")
    if a in ("click", "long_press", "swipe", "type", "open", "wait"):
        expected = {"action"}
        if a in ("click", "long_press", "swipe"):
            expected.add("coordinate")
        if a == "swipe":
            expected.add("coordinate2")
        if a in ("type", "open"):
            expected.add("text")
        if set(action) != expected:
            raise ValueError("mobile105_unrepresentable_parameters:" + str(a))
        return dict(
            action, action={"click": "tap", "type": "type_text", "open": "open_app"}.get(a, a)
        )
    if a == "scroll" and set(action) == {"action", "direction"}:
        name = "scroll_" + str(action["direction"])
        if name in ACTIONS:
            return {"action": name}
    if a == "system_button" and set(action) == {"action", "button"}:
        names = {"Back": "back", "Home": "home", "Enter": "enter", "Menu": "menu"}
        if action["button"] in names:
            return {"action": names[action["button"]]}
    if a == "key" and set(action) == {"action", "text"}:
        names = {"BACK": "back", "HOME": "home", "ENTER": "enter", "MENU": "menu"}
        if action["text"] in names:
            return {"action": names[action["text"]]}
    if (
        a == "terminate"
        and set(action) == {"action", "status"}
        and action["status"] in ("success", "failure")
    ):
        return {"action": "done" if action["status"] == "success" else "failed"}
    raise ValueError("unsupported_mobile105_source_action:" + str(action))


def to_source_action(action):
    a = action.get("action")
    if a not in ACTIONS:
        raise ValueError("unknown_mobile105_action")
    if a in ("tap", "long_press", "swipe", "type_text", "open_app", "wait"):
        result = dict(
            action, action={"tap": "click", "type_text": "type", "open_app": "open"}.get(a, a)
        )
    elif set(action) != {"action"}:
        raise ValueError("mobile105_unexpected_parameters")
    elif a.startswith("scroll_"):
        result = {"action": "scroll", "direction": a.removeprefix("scroll_")}
    elif a in ("back", "home", "enter", "menu"):
        result = {"action": "system_button", "button": a.capitalize()}
    else:
        result = {"action": "terminate", "status": "success" if a == "done" else "failure"}
    # Reject time or any silently ignored extra fields, including for history.
    if from_source_action(result) != action:
        raise ValueError("mobile105_noncanonical_action")
    return result


@dataclass(frozen=True)
class Mobile105ExecutionConfig:
    protocol: str = PROTOCOL
    primitives: MobileExecutionConfig = field(default_factory=MobileExecutionConfig)

    def __post_init__(self):
        if self.protocol != PROTOCOL or not isinstance(self.primitives, MobileExecutionConfig):
            raise ValueError("mobile105_execution_contract_mismatch")


class Mobile105Codec:
    variant = PROTOCOL
    output_slots = 105

    def __init__(self, tokenizer, *, bundle=BUNDLE):
        self.source_codec = Mobile133Codec(tokenizer, bundle=bundle)
        self.upstream = self.source_codec.upstream
        for name in (
            "tokenizer",
            "base_codec",
            "bin_ids",
            "codebook_sha256",
            "mask_id",
            "pad_id",
            "eos_id",
            "special_ids",
            "vocab_size",
        ):
            setattr(self, name, getattr(self.source_codec, name))
        self.action_ids = {
            a: self.base_codec._single(label)
            for label, a in zip(ACTION_LABELS, ACTIONS, strict=True)
        }
        if len(set(self.action_ids.values())) != len(ACTIONS):
            raise ValueError("mobile105_action_code_collision")
        built = self.build_parallel_input()
        self.input_ids, self.slot_positions = built.input_ids, built.slot_positions
        self.field_slices = {}
        start = 0
        for name, width in FIELD_WIDTHS:
            self.field_slices[name] = slice(start, start + width)
            start += width
        assert start == self.output_slots

    def build_parallel_input(self, prefix_ids=()):
        from trace2task.gui_owl_short_parallel_input import ParallelInput

        ids, positions, fields = list(prefix_ids), [], []
        for literal, (name, width) in zip(LITERALS[:-1], FIELD_WIDTHS, strict=True):
            ids.extend(self.tokenizer.encode(literal, add_special_tokens=False).ids)
            pos = tuple(range(len(ids), len(ids) + width))
            positions.extend(pos)
            fields.append((name, pos))
            ids.extend([self.mask_id] * width)
        ids.extend(self.tokenizer.encode(LITERALS[-1], add_special_tokens=False).ids)
        return ParallelInput(tuple(ids), tuple(positions), tuple(fields), len(prefix_ids))

    def contract(self):
        return {
            "protocol": PROTOCOL,
            "output_slots": 105,
            "layout": [[name, width] for name, width in FIELD_WIDTHS],
            "action_codes": [
                {"action": a, "label": label, "token_id": self.action_ids[a]}
                for label, a in zip(ACTION_LABELS, ACTIONS, strict=True)
            ],
            "coordinate_codebook_sha256": self.codebook_sha256,
            "scaffold_sha256": hashlib.sha256(json.dumps(self.input_ids).encode()).hexdigest(),
            "max_text_tokens": 99,
            "text_slots_including_eos": 100,
            "timing": "executor defaults, no duration labels",
            "scroll_direction": "browsing direction; no predicted endpoints",
            "inactive_slots": "ignored in supervision and execution",
            "readout": "native causal logits at slot_position-1; single forward",
        }

    def validate_contract(self, value):
        if value != self.contract():
            raise ValueError("mobile105_protocol_contract_mismatch")

    def adapt_target(self, action, *, coordinate_supervision=None):
        source = to_source_action(action)
        t = self.source_codec.adapt_target(
            source, "android", coordinate_supervision=coordinate_supervision
        )
        canonical = from_source_action(t.canonical_action)
        ids, hard = [self.pad_id] * 105, [-100] * 105
        ids[0] = hard[0] = self.action_ids[canonical["action"]]
        soft, fields = {}, [("action", 0, 1)]
        for name, start, stop in t.fields:
            if name not in self.field_slices or name == "action":
                continue  # direction/button/status are encoded by the action class.
            for old in range(start, stop):
                new = PARAMETER_POSITIONS[old]
                ids[new], hard[new] = t.target_ids[old], t.hard_labels[old]
                if old in t.soft_labels:
                    soft[new] = t.soft_labels[old]
            fields.append((name, PARAMETER_POSITIONS[start], PARAMETER_POSITIONS[stop - 1] + 1))
        return replace(
            t,
            canonical_action=canonical,
            target_ids=tuple(ids),
            hard_labels=tuple(hard),
            soft_labels=soft,
            fields=tuple(fields),
        )

    def adapt_training_target(self, row):
        self.source_codec.check_row(row)
        if row["platform"] != "android":
            raise ValueError("mobile105_requires_android")
        return self.adapt_target(
            from_source_action(row["target"]),
            coordinate_supervision=self.upstream.coordinate_metadata_from_row(row),
        )

    def messages(self, row):
        self.source_codec.check_row(row)
        if row["platform"] != "android":
            raise ValueError("mobile105_requires_android")
        history = [
            dict(step, actions=[from_source_action(a) for a in step["actions"]])
            for step in row["history"]
        ]
        # Reuse the pinned image order, complete-history formatter and escaping.
        result = self.upstream.messages(dict(row, history=history))
        result[0]["content"] = SYSTEM
        result[1]["content"][0]["text"] = result[1]["content"][0]["text"].replace(
            "\nTask: ", "\nCurrent instruction: ", 1
        )
        return result

    def bind_history_tokens(self, prefix_ids, row):
        # The action-name projection preserves coordinate order and every history item.
        return self.source_codec.bind_history_tokens(prefix_ids, row)

    def decode(self, logits):
        import torch

        if logits.ndim != 2 or logits.shape[0] != 105 or logits.shape[1] < self.vocab_size:
            raise ValueError("expected_105_by_vocab_logits")
        if not torch.isfinite(logits).all():
            raise ValueError("nonfinite_logits")
        a = ACTIONS[int(logits[0, list(self.action_ids.values())].argmax())]
        if a not in ("tap", "long_press", "swipe", "type_text", "open_app"):
            return {"action": a}
        # Parameter-only compatibility adapter: no second model call, no new readout rule.
        full = logits.new_full((133, logits.shape[1]), -60)
        old = list(PARAMETER_POSITIONS)
        full[old] = logits[list(PARAMETER_POSITIONS.values())]
        old_name = {"tap": "click", "type_text": "type", "open_app": "open"}.get(a, a)
        for p, token in enumerate(self.source_codec.action_encodings[old_name]):
            full[p, token] = 60
        if a == "long_press":
            full[self.source_codec.field_slices["time"], self.pad_id] = 60
        decoded = self.source_codec.decode(full, platform="android")
        return from_source_action(decoded)

    def execution_plan(self, action, size, *, config, installed_apps=None):
        if not isinstance(config, Mobile105ExecutionConfig) or config.protocol != PROTOCOL:
            raise ValueError("mobile105_execution_contract_mismatch")
        canonical = self.adapt_target(action).canonical_action
        return self.source_codec.execution_plan(
            to_source_action(canonical),
            size,
            config=config.primitives,
            installed_apps=installed_apps,
        )