registry.py

"""Runtime Registry — single source of truth for all live agent state.

Seeded from config.toml base definitions. All systems read from the registry
at runtime, never from config.toml directly.

Thread-safe: a single threading.Lock guards all mutations.
"""

import colorsys
import json
import secrets
import threading
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path


@dataclass
class Instance:
    """A live agent instance."""
    name: str       # canonical ID: "gemini", "gemini-2"
    base: str       # base family: "gemini"
    slot: int       # 1, 2, 3...
    label: str      # "Gemini", "Gemini 2", or human-set custom
    color: str      # hex color (derived from base + slot)
    identity_id: str = field(default_factory=lambda: uuid.uuid4().hex)
    token: str = field(default_factory=lambda: secrets.token_hex(16))
    epoch: int = 1
    state: str = "pending"   # "pending" | "active"
    registered_at: float = field(default_factory=time.time)


class RuntimeRegistry:
    GRACE_PERIOD = 30  # seconds — name reserved after deregister

    def __init__(self, data_dir: str = "./data"):
        self._lock = threading.Lock()
        self._bases: dict[str, dict] = {}          # base name → config template
        self._instances: dict[str, Instance] = {}   # canonical name → Instance
        self._reserved: dict[str, float] = {}       # name → deregister timestamp
        self._renames: dict[str, str] = {}           # old name → new name (for heartbeat redirect)
        self._on_change_cbs: list = []
        self._data_dir = Path(data_dir)
        self._load_renames()

    # --- Setup ---

    def seed(self, agents_config: dict):
        """Load base templates from config.toml [agents.*] section."""
        with self._lock:
            for name, cfg in agents_config.items():
                self._bases[name] = dict(cfg)

    def on_change(self, cb):
        """Register a callback fired after any registry mutation."""
        self._on_change_cbs.append(cb)

    def _notify(self):
        for cb in self._on_change_cbs:
            try:
                cb()
            except Exception:
                pass

    # --- Rename persistence ---

    def _renames_path(self) -> Path:
        return self._data_dir / "renames.json"

    def _load_renames(self):
        p = self._renames_path()
        if p.exists():
            try:
                self._renames = json.loads(p.read_text("utf-8"))
            except Exception:
                self._renames = {}

    def _save_renames(self):
        """Persist renames to disk. Must be called outside the lock."""
        try:
            self._data_dir.mkdir(parents=True, exist_ok=True)
            tmp = self._renames_path().with_suffix(".tmp")
            with self._lock:
                data = dict(self._renames)
            tmp.write_text(json.dumps(data), "utf-8")
            tmp.replace(self._renames_path())
        except Exception:
            pass

    # --- Registration ---

    def register(self, base: str, label: str | None = None) -> dict | None:
        """Register a new instance of `base`. Returns slot info or None if unknown base.

        When a 2nd instance registers, slot 1 is renamed from 'base' to 'base-1'
        to prevent identity ambiguity. The rename info is returned as '_renamed_slot1'.
        """
        with self._lock:
            if base not in self._bases:
                return None

            self._expire_reserved()

            # Find next free slot
            taken = {i.slot for i in self._instances.values() if i.base == base}
            reserved = set()
            for rn in self._reserved:
                rb, rs = self._parse_name(rn)
                if rb == base:
                    reserved.add(rs)

            slot = 1
            while slot in taken or slot in reserved:
                slot += 1

            # When a 2nd instance registers, rename slot-1 from "base" to "base-1"
            # so that no instance shares a name with the base family.  This prevents
            # a second instance from sending messages as "base" (identity theft).
            renamed_slot1 = None
            if slot >= 2 and base in self._instances:
                slot1 = self._instances[base]
                if slot1.base == base and slot1.slot == 1:
                    new_s1_name = f"{base}-1"
                    del self._instances[base]
                    slot1.name = new_s1_name
                    base_cfg = self._bases[base]
                    slot1.label = f"{base_cfg.get('label', base.capitalize())} 1"
                    # Color stays the same (slot 1 = base color)
                    self._instances[new_s1_name] = slot1
                    self._renames[base] = new_s1_name
                    renamed_slot1 = {"old": base, "new": new_s1_name}

            name = base if slot == 1 else f"{base}-{slot}"
            base_cfg = self._bases[base]
            color = _derive_color(base_cfg.get("color", "#888"), slot)

            if label:
                lbl = label
            elif slot == 1:
                lbl = base_cfg.get("label", base.capitalize())
            else:
                lbl = f"{base_cfg.get('label', base.capitalize())} {slot}"

            # Fresh registrations are immediately authoritative. Identity
            # recovery/reclaim still uses chat_claim, but normal startup should
            # not block on a manual confirmation step.
            state = "active"
            inst = Instance(name=name, base=base, slot=slot, label=lbl, color=color, state=state)
            self._instances[name] = inst
            result = _inst_dict(inst, include_token=True)
            if renamed_slot1:
                result["_renamed_slot1"] = renamed_slot1

        self._notify()
        self._save_renames()
        return result

    def deregister(self, name: str) -> dict | None:
        """Remove an instance. Name is reserved for GRACE_PERIOD seconds.

        Returns result dict with 'ok' and optional '_renamed_back' info,
        or None if instance not found.
        """
        with self._lock:
            if name not in self._instances:
                return None
            base = self._instances[name].base
            del self._instances[name]
            self._reserved[name] = time.time()

            # If family drops to 1 instance with a numbered name, rename back to base
            renamed_back = None
            family = [i for i in self._instances.values() if i.base == base]
            if len(family) == 1:
                remaining = family[0]
                r_base, r_slot = self._parse_name(remaining.name)
                if r_base == base and remaining.name != base:
                    old_name = remaining.name
                    del self._instances[old_name]
                    remaining.name = base
                    remaining.slot = 1
                    base_cfg = self._bases.get(base, {})
                    remaining.label = base_cfg.get("label", base.capitalize())
                    remaining.color = _derive_color(base_cfg.get("color", "#888"), 1)
                    self._instances[base] = remaining
                    self._renames[old_name] = base
                    renamed_back = {"old": old_name, "new": base}

        self._notify()
        self._save_renames()
        result = {"ok": True}
        if renamed_back:
            result["_renamed_back"] = renamed_back
        return result

    # --- Identity Claim ---

    def claim(self, sender: str, target_name: str | None = None) -> dict | str:
        """Claim an identity. Returns instance dict on success, error string on failure.

        Family-based matching: sender can be a base family name (e.g. 'claude')
        and the server assigns the next unclaimed instance of that family.

        - sender='claude', no target: assign first unclaimed claude instance
        - sender='claude', target='claude-music': assign unclaimed instance AND rename
        - sender='claude-2' (exact match): confirm that specific instance
        """
        error = None
        result = None
        with self._lock:
            inst = None

            # If sender is a base family name, use family-based matching
            # (don't exact-match the slot-1 instance — that causes both
            # callers to claim the same identity)
            if sender in self._bases:
                # Find first unclaimed (pending) instance in this family
                for candidate in self._instances.values():
                    if candidate.base == sender and candidate.state == "pending":
                        inst = candidate
                        break
                # If no pending, fall back to any instance in the family
                if not inst:
                    for candidate in self._instances.values():
                        if candidate.base == sender:
                            inst = candidate
                            break
            else:
                # Exact match for specific instance names (e.g. 'claude-2')
                inst = self._instances.get(sender)

            if not inst:
                error = f"No available {sender} instance. Is a wrapper registered?"
            elif target_name is None or target_name == inst.name:
                # Accept current name — but don't auto-activate pending instances.
                # Pending instances must be named by human (lightbox) or reclaimed
                # with an explicit target name (breadcrumb resume).
                if inst.state == "pending" and target_name is None:
                    result = _inst_dict(inst)  # return info but stay pending
                else:
                    inst.state = "active"
                    result = _inst_dict(inst)
            else:
                # Rename/reclaim — check collision and family guard
                if target_name in self._instances and target_name != inst.name:
                    error = f"Already claimed: {target_name}"
                elif (family_err := self._conflicts_with_other_family(target_name, inst.base)):
                    error = family_err
                else:
                    # Check slot collision within same family
                    t_base, t_slot = self._parse_name(target_name)
                    if t_base == inst.base:
                        slot_taken = any(
                            i.slot == t_slot and i.name != inst.name
                            for i in self._instances.values() if i.base == inst.base
                        )
                        if slot_taken:
                            error = f"Slot {t_slot} already occupied in {inst.base} family"
                    if not error:
                        # Swap identity to target name
                        self._reserved.pop(target_name, None)
                        old_name = inst.name
                        del self._instances[old_name]
                        inst.name = target_name
                        inst.state = "active"
                        # Recalculate slot, color, and label from target name
                        base_cfg = self._bases.get(inst.base, {})
                        if t_base == inst.base:
                            # Target parses as same family (e.g. 'claude' or 'claude-3')
                            inst.slot = t_slot
                            inst.color = _derive_color(base_cfg.get("color", "#888"), t_slot)
                            if t_slot == 1:
                                inst.label = base_cfg.get("label", inst.base.capitalize())
                            else:
                                inst.label = f"{base_cfg.get('label', inst.base.capitalize())} {t_slot}"
                        else:
                            # Custom name (e.g. 'claude-music') — keep slot color, use name as label
                            inst.label = target_name
                        self._instances[target_name] = inst
                        # Track rename so wrapper can discover it via heartbeat
                        self._renames[old_name] = target_name
                        result = _inst_dict(inst)

        if error:
            return error
        self._notify()
        self._save_renames()
        return result

    def confirm_pending(self, name: str) -> bool:
        """Auto-confirm a pending instance (10s timeout path)."""
        with self._lock:
            inst = self._instances.get(name)
            if not inst or inst.state != "pending":
                return False
            inst.state = "active"

        self._notify()
        self._save_renames()
        return True

    # --- Rename / Label ---

    def rename(self, old_name: str, new_name: str, label: str | None = None) -> dict | str:
        """Full identity rename (human-initiated). Returns instance dict or error string.

        Changes the sender ID, label, and tracks the rename for wrapper sync.
        If new_name == old_name, falls back to a label-only change.
        """
        with self._lock:
            inst = self._instances.get(old_name)
            if not inst:
                return f"Not found: {old_name}"

            if new_name == old_name:
                # Same identity — just update label
                if label:
                    inst.label = label
                result = _inst_dict(inst)
            elif new_name in self._instances:
                return f"Already taken: {new_name}"
            elif (family_err := self._conflicts_with_other_family(new_name, inst.base)):
                return family_err
            else:
                # Check slot collision within same family
                t_base, t_slot = self._parse_name(new_name)
                if t_base == inst.base:
                    slot_taken = any(
                        i.slot == t_slot and i.name != old_name
                        for i in self._instances.values() if i.base == inst.base
                    )
                    if slot_taken:
                        return f"Slot {t_slot} already occupied in {inst.base} family"

                # Move instance to new name
                del self._instances[old_name]
                inst.name = new_name

                # Set label (use provided label, or derive from new_name)
                base_cfg = self._bases.get(inst.base, {})
                if label:
                    inst.label = label
                elif t_base == inst.base and t_slot != inst.slot:
                    # Numbered variant (e.g. claude-3) — use "Claude 3"
                    if t_slot == 1:
                        inst.label = base_cfg.get("label", inst.base.capitalize())
                    else:
                        inst.label = f"{base_cfg.get('label', inst.base.capitalize())} {t_slot}"
                else:
                    inst.label = new_name

                # Update slot + color if it's a numbered family name
                if t_base == inst.base:
                    inst.slot = t_slot
                    inst.color = _derive_color(base_cfg.get("color", "#888"), t_slot)

                self._instances[new_name] = inst
                self._renames[old_name] = new_name
                result = _inst_dict(inst)

        self._notify()
        self._save_renames()
        return result

    def set_label(self, name: str, label: str) -> bool:
        """Set display label only (no identity change)."""
        with self._lock:
            inst = self._instances.get(name)
            if not inst:
                return False
            inst.label = label

        self._notify()
        self._save_renames()
        return True

    # --- Queries ---

    def get_instance(self, name: str) -> dict | None:
        with self._lock:
            inst = self._instances.get(name)
            return _inst_dict(inst) if inst else None

    def get_all(self) -> dict[str, dict]:
        """All registered instances as {name: {name, base, slot, label, color, state}}."""
        with self._lock:
            return {n: _inst_dict(i) for n, i in self._instances.items()}

    def get_agent_config(self) -> dict[str, dict]:
        """For WebSocket 'agents' message: {name: {color, label, base, state}}."""
        with self._lock:
            return {
                n: {"color": i.color, "label": i.label, "base": i.base, "state": i.state}
                for n, i in self._instances.items()
            }

    def get_all_names(self) -> list[str]:
        with self._lock:
            return list(self._instances.keys())

    def get_active_names(self) -> list[str]:
        with self._lock:
            return [n for n, i in self._instances.items() if i.state == "active"]

    def get_instances_for(self, base: str) -> list[dict]:
        with self._lock:
            return [_inst_dict(i) for i in self._instances.values() if i.base == base]

    def get_bases(self) -> dict[str, dict]:
        with self._lock:
            return dict(self._bases)

    def get_base_config(self, base: str) -> dict | None:
        with self._lock:
            return dict(self._bases[base]) if base in self._bases else None

    def is_agent_family(self, name: str) -> bool:
        """Check if a name belongs to any agent family (base, slot, or custom alias)."""
        with self._lock:
            # Check registered instance first (handles custom names like 'claude-music')
            inst = self._instances.get(name)
            if inst:
                return inst.base in self._bases
            # Fall back to name parsing for slot names like 'claude-2'
            base, _ = self._parse_name(name)
            if base in self._bases:
                return True
            # Treat unregistered custom aliases like 'claude-prime' as belonging
            # to the same family so stale senders are rejected until claimed.
            return any(name.startswith(f"{family}-") for family in self._bases)

    def family_instance_count(self, name: str) -> int:
        """Count registered instances in the same family as `name`."""
        with self._lock:
            # Check registered instance first (handles custom names)
            inst = self._instances.get(name)
            if inst:
                base = inst.base
            else:
                base, _ = self._parse_name(name)
                if base not in self._bases:
                    for family in self._bases:
                        if name.startswith(f"{family}-"):
                            base = family
                            break
            return sum(1 for i in self._instances.values() if i.base == base)

    def has_claimed_instances(self, base: str) -> bool:
        """Check if any instance in this family has been claimed (state=active)."""
        with self._lock:
            return any(
                i.state == "active" and i.base == base
                for i in self._instances.values()
            )

    def get_family_instance(self, base: str) -> dict | None:
        """Return the instance dict for a family if exactly one exists.
        Used by heartbeat to find renamed instances after server restart."""
        with self._lock:
            matches = [i for i in self._instances.values() if i.base == base]
            if len(matches) == 1:
                return _inst_dict(matches[0])
        return None

    def resolve_to_instances(self, name: str) -> list[str]:
        """Resolve a name to actual registered instance names.

        If `name` is a registered instance, returns [name].
        If `name` is a base family name with no exact match, returns all
        active instances in that family (e.g. 'claude' → ['claude-prime']).
        Otherwise returns [name] unchanged (for non-agent names like 'ben').
        """
        with self._lock:
            if name in self._instances:
                return [name]
            # Check if it's a base name with registered family members
            if name in self._bases:
                members = [i.name for i in self._instances.values()
                           if i.base == name and i.state == "active"]
                if members:
                    return members
            return [name]

    def resolve_name(self, name: str) -> str:
        """Follow rename chain to find current canonical name."""
        with self._lock:
            # Follow renames (e.g. claude-2 → claude-music)
            seen = set()
            current = name
            while current in self._renames and current not in seen:
                seen.add(current)
                current = self._renames[current]
            return current

    def is_registered(self, name: str) -> bool:
        with self._lock:
            return name in self._instances

    def is_pending(self, name: str) -> bool:
        with self._lock:
            i = self._instances.get(name)
            return i is not None and i.state == "pending"

    def resolve_token(self, token: str) -> dict | None:
        """Map an instance_token to the current canonical instance dict, or None."""
        with self._lock:
            for inst in self._instances.values():
                if inst.token == token:
                    return _inst_dict(inst)
        return None

    def get_pending(self) -> list[dict]:
        """All pending instances (for timeout checks)."""
        with self._lock:
            return [_inst_dict(i) for i in self._instances.values()
                    if i.state == "pending"]

    # --- Internal ---

    def _conflicts_with_other_family(self, name: str, own_base: str) -> str | None:
        """Check if `name` stomps on another family's namespace.

        Returns an error string if it conflicts, None if safe.
        Blocks: renaming claude to 'gemini', 'gemini-2', 'codex', etc.
        Allows: renaming claude to 'cudders', 'claude-prime', etc.
        """
        t_base, _ = self._parse_name(name)
        # If the parsed base matches a known family that isn't ours, block it
        if t_base in self._bases and t_base != own_base:
            return f"Name '{name}' conflicts with the {t_base} agent family"
        # Also block if the raw name exactly matches another family's base
        if name in self._bases and name != own_base:
            return f"Name '{name}' is a reserved agent family name"
        return None

    def _parse_name(self, name: str) -> tuple[str, int]:
        """Parse 'gemini-2' -> ('gemini', 2), 'gemini' -> ('gemini', 1)."""
        if "-" in name:
            prefix, suffix = name.rsplit("-", 1)
            try:
                return prefix, int(suffix)
            except ValueError:
                pass
        return name, 1

    def clean_renames_for(self, name: str):
        """Remove all rename chain entries pointing to or from `name`."""
        with self._lock:
            # Remove entries where name is a key (old name → ...)
            self._renames.pop(name, None)
            # Remove entries where name is a value (... → name)
            stale = [k for k, v in self._renames.items() if v == name]
            for k in stale:
                del self._renames[k]
        self._save_renames()

    def _expire_reserved(self):
        """Remove expired reservations. Must hold lock."""
        now = time.time()
        self._reserved = {n: t for n, t in self._reserved.items()
                          if now - t < self.GRACE_PERIOD}


# --- Module-level helpers ---

def _inst_dict(inst: Instance, include_token: bool = False) -> dict:
    d = {
        "identity_id": inst.identity_id,
        "name": inst.name, "base": inst.base, "slot": inst.slot,
        "label": inst.label, "color": inst.color, "state": inst.state,
        "epoch": inst.epoch,
        "registered_at": inst.registered_at,
    }
    if include_token:
        d["token"] = inst.token
    return d


def _derive_color(base_hex: str, slot: int) -> str:
    """Derive variant color: slot 1 = base, slot N = hue/lightness shifted.

    Pattern: slot 2 = hue +25 deg, L +5%; slot 3 = hue -25 deg, L -5%; etc.
    """
    if slot == 1:
        return base_hex
    hx = base_hex.lstrip("#")
    if len(hx) != 6:
        return base_hex
    r, g, b = int(hx[0:2], 16), int(hx[2:4], 16), int(hx[4:6], 16)
    h, l, s = colorsys.rgb_to_hls(r / 255, g / 255, b / 255)

    # Alternating hue shifts with increasing magnitude
    magnitude = ((slot - 1 + 1) // 2) * 25
    direction = 1 if slot % 2 == 0 else -1
    h = (h + direction * magnitude / 360) % 1.0
    l = max(0.15, min(0.85, l + direction * 0.05))

    r2, g2, b2 = colorsys.hls_to_rgb(h, l, s)
    return f"#{int(r2 * 255):02x}{int(g2 * 255):02x}{int(b2 * 255):02x}"