"""Synthetic dataset for capability C (the SIGIL protocol).

A few hundred (request -> directive) pairs, generated deterministically from a
fixed seed so a reviewer regenerates byte-identical data. The held-out eval set
uses locations and numeric values that never appear in training, so a high eval
score means the model learned the protocol and generalises the copy slot, not
that it memorised specific strings.
"""

from __future__ import annotations

import random
from dataclasses import dataclass
from typing import List, Tuple

from .capability import Directive

# Domain nouns -> route key. Several surface nouns per domain so phrasing varies.
_DOMAIN_NOUNS = {
    "lights": ["lights", "lamp", "lighting"],
    "door": ["door", "doorway"],
    "thermostat": ["thermostat", "heating"],
    "alarm": ["alarm", "security alarm"],
    "speaker": ["speaker", "music"],
    "blinds": ["blinds", "shades"],
}

# Action phrasings -> verb key.
_ACTION_PHRASES = {
    "on": ["turn on", "switch on", "activate", "power up", "enable"],
    "off": ["turn off", "switch off", "deactivate", "power down", "disable"],
    "query": ["check", "what's the status of", "get the status of", "report on"],
    "set": ["set", "adjust", "change"],
    "cancel": ["cancel the schedule for", "clear the timer on", "remove the automation on"],
}

# Locations: disjoint train / eval so the copy slot is tested on unseen values.
_TRAIN_LOCATIONS = [
    "kitchen", "bedroom", "garage", "hallway", "office", "bathroom",
    "basement", "porch", "attic", "pantry", "nursery", "workshop",
    "cellar", "foyer", "landing", "study", "utility", "garden",
    "balcony", "corridor", "stairwell", "cloakroom", "larder", "parlour",
    "sunroom", "playroom", "gym", "library", "den",
    "terrace", "courtyard", "driveway", "shed", "greenhouse", "laundry",
]
_EVAL_LOCATIONS = [
    "lounge", "veranda", "mudroom", "scullery", "snug", "loft",
    "vestibule", "annexe",
]

# BIND (set) applies only to domains that take a numeric value.
_VALUE_DOMAINS = ["thermostat", "speaker", "blinds"]
_TRAIN_VALUES = [17, 18, 19, 20, 21, 22, 30, 40, 50, 60]
_EVAL_VALUES = [16, 23, 24, 45, 70]


@dataclass(frozen=True)
class Example:
    request: str
    directive: Directive


def _location_example(rng: random.Random, action: str, locations: List[str]) -> Example:
    domain = rng.choice(list(_DOMAIN_NOUNS))
    # BANISH/query/on/off all take a location arg.
    noun = rng.choice(_DOMAIN_NOUNS[domain])
    location = rng.choice(locations)
    phrase = rng.choice(_ACTION_PHRASES[action])
    request = f"{phrase.capitalize()} the {location} {noun}.".replace("..", ".")
    directive = Directive(
        route=_route_for(domain), verb=_verb_for(action), arg=location
    )
    return Example(request=request, directive=directive)


def _value_example(rng: random.Random, values: List[int]) -> Example:
    domain = rng.choice(_VALUE_DOMAINS)
    noun = rng.choice(_DOMAIN_NOUNS[domain])
    value = rng.choice(values)
    phrase = rng.choice(_ACTION_PHRASES["set"])
    request = f"{phrase.capitalize()} the {noun} to {value}."
    directive = Directive(route=_route_for(domain), verb="BIND", arg=str(value))
    return Example(request=request, directive=directive)


def _route_for(domain: str) -> str:
    from .capability import ROUTES

    return ROUTES[domain]


def _verb_for(action: str) -> str:
    from .capability import VERBS

    return VERBS[action]


def _build_split(rng: random.Random, n: int, locations: List[str], values: List[int]) -> List[Example]:
    examples: List[Example] = []
    location_actions = ["on", "off", "query", "cancel"]
    for _ in range(n):
        # ~4/5 location-arg examples, ~1/5 value-arg (BIND) examples, balanced.
        if rng.random() < 0.75:
            action = rng.choice(location_actions)
            examples.append(_location_example(rng, action, locations))
        else:
            examples.append(_value_example(rng, values))
    return examples


def build_dataset(
    n_train: int = 560, n_eval: int = 80, seed: int = 1729
) -> Tuple[List[Example], List[Example]]:
    """Return (train, eval). Eval args are disjoint from train args."""
    rng = random.Random(seed)
    train = _build_split(rng, n_train, _TRAIN_LOCATIONS, _TRAIN_VALUES)
    eval_rng = random.Random(seed + 1)
    held_out = _build_split(eval_rng, n_eval, _EVAL_LOCATIONS, _EVAL_VALUES)
    return train, held_out