"""Agentic MCP tools held *inside* HDC storage.

A tool is a record, not a row in a database:

    record = bundle([
        bind(role("name"),  hvc("tool/hdc.bind")),
        bind(role("cap"),   capability_vector("hdc/bind/construct")),
        bind(role("arity"), hvc("arity/2")),
    ])

The record is written to storage as the cell ``toolreg/<name>``.  Registration,
discovery and dispatch therefore all run on the same content-addressable
machinery as programs and data: a tool is found by *resembling* the query.

Dispatch has three tiers, each reporting its confidence:

1. exact name — confidence 1.0;
2. character-superposition similarity, so a misspelled or truncated tool name
   (``hdc.bundl``) still finds ``hdc.bundle`` rather than failing;
3. capability path similarity, so ``hdc.memory.search`` reaches
   ``hdc.recall`` because both carry ``hdc`` + ``memory`` atoms.

Tier 2 uses the standard HDC string codec: each character becomes an atom and
its position is bound in by permutation, so two strings that match in most
positions have highly similar vectors.  That is what makes fuzzy dispatch a
cleanup operation instead of an edit-distance algorithm.

The registry is what the ``TOOL`` instruction calls, and what
:mod:`hdc.mcp_server` advertises to an outside agent, so a program running
inside HDC and an MCP client reach the same tools by the same path.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable, Sequence

from .hypervec import (
    DEFAULT_DIM,
    bind,
    bundle,
    hvc,
    hv_to_int,
    int_to_hv,
    permute,
    popcount,
    sim,
)
from .memory import HDCStorage

#: Minimum fuzzy similarity accepted for a non-exact tool name.
FUZZY_FLOOR = 0.55
#: Fuzzy matches must beat the runner-up by this much to be trusted.
FUZZY_MARGIN = 0.04

_CAP_CACHE: dict[tuple[str, int], int] = {}
_STR_CACHE: dict[tuple[str, int], int] = {}


def capability_vector(capability: str, d: int = DEFAULT_DIM) -> int:
    """Superpose the path atoms of a capability (``hdc/bind/construct``)."""
    key = (capability, d)
    hit = _CAP_CACHE.get(key)
    if hit is not None:
        return hit
    atoms = [hvc("cap/" + part, d) for part in capability.replace(".", "/").split("/") if part]
    vec = bundle(atoms, d) if atoms else hvc("cap/unknown", d)
    _CAP_CACHE[key] = vec
    return vec


def string_vector(text: str, d: int = DEFAULT_DIM) -> int:
    """Character-superposition codec: order-sensitive, near-orthogonal, fuzzy.

    Each character is bound to its position by permutation and the results are
    bundled, so ``sim(string_vector("hdc.bundle"), string_vector("hdc.bundl"))``
    is high while unrelated strings stay near the noise floor.
    """
    key = (text, d)
    hit = _STR_CACHE.get(key)
    if hit is not None:
        return hit
    if not text:
        vec = 0
    else:
        vec = bundle(
            [permute(hvc("char/" + ch, d), i, d) for i, ch in enumerate(text)], d
        )
    _STR_CACHE[key] = vec
    return vec


def _json_type(kind: str) -> str:
    return "integer" if kind in ("int", "hv") else "string"


@dataclass
class ToolSpec:
    """One agentic tool.

    ``returns`` matters because a hypervector and a scalar are both Python
    ``int``s: the machine has to know whether to *encode* a result (scalar) or
    store it as a vector in its own right.  Without this the two are
    indistinguishable at the call site.
    """

    name: str
    capability: str
    params: tuple[tuple[str, str], ...]
    doc: str
    fn: Callable[..., object]
    returns: str = "int"
    calls: int = 0

    @property
    def kinds(self) -> tuple[str, ...]:
        return tuple(kind for _, kind in self.params)

    def record(self, d: int) -> int:
        """The hypervector record that represents this tool in storage."""
        parts = [
            bind(self._role("name", d), hvc("tool/" + self.name, d), d),
            bind(self._role("cap", d), capability_vector(self.capability, d), d),
            bind(self._role("arity", d), hvc("arity/%d" % len(self.params), d), d),
        ]
        return bundle(parts, d)

    @staticmethod
    def _role(name: str, d: int) -> int:
        from .hypervec import role

        return role(name, d)

    def input_schema(self) -> dict:
        """MCP ``inputSchema`` built from the declared parameter kinds."""
        props = {}
        for pname, kind in self.params:
            props[pname] = {
                "type": _json_type(kind),
                "description": (
                    "hypervector: pass the integer, or a hex string like 0x1f3a"
                    if kind == "hv"
                    else "scalar integer"
                ),
            }
        return {
            "type": "object",
            "properties": props,
            "required": [pname for pname, _ in self.params],
            "additionalProperties": False,
        }


class ToolRegistry:
    """Tool records in HDC storage, plus the dispatcher that reads them."""

    def __init__(self, storage: HDCStorage | None = None, d: int = DEFAULT_DIM):
        self.d = d
        self.storage = storage if storage is not None else HDCStorage(d=d, label="tools")
        self.specs: dict[str, ToolSpec] = {}
        self._records: dict[str, int] = {}
        self.last_machine = None
        self.install_builtins()

    # ---------------------------------------------------------- registration
    def register(
        self,
        name: str,
        capability: str,
        params: Sequence[tuple[str, str]],
        doc: str,
        fn: Callable[..., object],
        returns: str = "int",
        store: bool = True,
    ) -> ToolSpec:
        """Add a tool and write its record into HDC storage."""
        spec = ToolSpec(name, capability, tuple(params), doc, fn, returns=returns)
        self.specs[name] = spec
        rec = spec.record(self.d)
        self._records[name] = rec
        if store:
            self.storage.write("toolreg/" + name, rec, symbol=True)
        return spec

    def unregister(self, name: str) -> None:
        self.specs.pop(name, None)
        self._records.pop(name, None)
        if "toolreg/" + name in self.storage:
            self.storage.cells.pop("toolreg/" + name, None)
            self.storage.order.remove("toolreg/" + name)

    # -------------------------------------------------------------- lookup
    def get(self, name: str) -> ToolSpec | None:
        return self.specs.get(name)

    def params_for(self, name: str) -> tuple[str, ...]:
        spec, _ = self.resolve(name)
        if spec is None:
            raise ToolError("unknown tool: %r" % (name,))
        return spec.kinds

    def returns_for(self, name: str) -> str:
        spec, _ = self.resolve(name)
        if spec is None:
            raise ToolError("unknown tool: %r" % (name,))
        return spec.returns

    def resolve(self, query: str | int) -> tuple[ToolSpec | None, float]:
        """Find a tool by name, by fuzzy name, or by capability resemblance."""
        if isinstance(query, int):
            best_name, best = "", -2.0
            for name, rec in self._records.items():
                s = sim(rec, query, self.d)
                if s > best:
                    best_name, best = name, s
            return (self.specs.get(best_name), best) if best >= 0.3 else (None, best)

        if query in self.specs:
            return self.specs[query], 1.0

        probe = string_vector(query, self.d)
        scored = sorted(
            ((name, sim(probe, string_vector(name, self.d), self.d)) for name in self.specs),
            key=lambda t: t[1],
            reverse=True,
        )
        if scored and scored[0][1] >= FUZZY_FLOOR:
            margin = scored[0][1] - (scored[1][1] if len(scored) > 1 else 0.0)
            if margin >= FUZZY_MARGIN or len(scored) == 1:
                return self.specs[scored[0][0]], scored[0][1]

        cap_probe = capability_vector(query, self.d)
        scored = sorted(
            (
                (name, sim(cap_probe, capability_vector(s.capability, self.d), self.d))
                for name, s in self.specs.items()
            ),
            key=lambda t: t[1],
            reverse=True,
        )
        if scored and scored[0][1] >= 0.30:
            return self.specs[scored[0][0]], scored[0][1]
        return None, (scored[0][1] if scored else 0.0)

    # --------------------------------------------------------------- calling
    def call(self, name: str, args: Sequence[object], machine=None) -> object:
        """Invoke a tool.  Raises ``ToolError`` for an unresolvable tool."""
        spec, confidence = self.resolve(name)
        if spec is None:
            raise ToolError(
                "no tool resembling %r (best similarity %.3f)" % (name, confidence)
            )
        if len(args) != len(spec.params):
            raise ToolError(
                "tool %s takes %d argument(s), got %d"
                % (spec.name, len(spec.params), len(args))
            )
        self.last_machine = machine
        spec.calls += 1
        return spec.fn(*args)

    def call_named(self, name: str, arguments: dict, machine=None) -> object:
        """Invoke a tool with named arguments (the MCP calling convention)."""
        spec, confidence = self.resolve(name)
        if spec is None:
            raise ToolError("unknown tool %r (%.3f)" % (name, confidence))
        args = []
        for pname, kind in spec.params:
            if pname not in arguments:
                raise ToolError("missing argument %r for %s" % (pname, spec.name))
            value = arguments[pname]
            if kind in ("int", "hv") and isinstance(value, str):
                value = int(value, 0)
            if kind == "int":
                value = int(value)
            elif kind == "hv":
                value = int(value)
            args.append(value)
        return self.call(spec.name, args, machine=machine)

    # ---------------------------------------------------------------- views
    def list_tools(self) -> list[dict]:
        """MCP ``tools/list`` payload."""
        out = []
        for name in sorted(self.specs):
            spec = self.specs[name]
            out.append(
                {
                    "name": spec.name,
                    "description": spec.doc,
                    "inputSchema": spec.input_schema(),
                    "annotations": {"hdc_capability": spec.capability},
                }
            )
        return out

    def specs_text(self) -> str:
        rows = ["%-16s %-26s %-5s %s" % ("TOOL", "CAPABILITY", "ARGS", "DESCRIPTION")]
        rows.append("-" * 96)
        for name in sorted(self.specs):
            s = self.specs[name]
            rows.append(
                "%-16s %-26s %-5d %s"
                % (s.name, s.capability, len(s.params), s.doc)
            )
        return "\n".join(rows)

    def memory_report(self) -> str:
        """Show the tool records as they sit in storage."""
        rows = ["tool records in %r (d=%d)" % (self.storage.label, self.d)]
        for name in sorted(self.specs):
            cell = "toolreg/" + name
            present = cell in self.storage
            hv = self.storage.read(cell, default=0)
            rows.append(
                "  %-16s cell=%-18s pop=%-5d %s"
                % (name, cell, popcount(hv), "stored" if present else "MISSING")
            )
        return "\n".join(rows)

    # ------------------------------------------------------------- builtins
    def install_builtins(self) -> None:
        """Register the built-in HDC / bit-logic / BitNet tools."""
        d = self.d
        store = self.storage

        def hdc_similarity(a: int, b: int) -> int:
            return int(round(sim(a, b, d) * 1000))

        def hdc_bind(a: int, b: int) -> int:
            return bind(a, b, d)

        def hdc_bundle(a: int, b: int) -> int:
            return bundle([a, b], d)

        def hdc_permute(v: int, k: int) -> int:
            return permute(v, k, d)

        def hdc_cleanup(v: int) -> int:
            hits = store.nearest_symbol(v, k=1)
            if not hits:
                return -1
            syms = store.symbols()
            return syms.index(hits[0][0]) if hits[0][0] in syms else -1

        def hdc_recall(v: int, k: int) -> int:
            hits = store.recall(v, k=max(1, k))
            if not hits:
                return -1
            return store.order.index(hits[0][0]) if hits[0][0] in store.order else -1

        def hdc_encode_int(n: int) -> int:
            return int_to_hv(n, d)

        def hdc_decode_int(v: int) -> int:
            return hv_to_int(v, d)

        def bit_popcount(v: int) -> int:
            return popcount(v)

        def bit_hamming(a: int, b: int) -> int:
            return (a ^ b).bit_count()

        def net_dot1bit(a: int, b: int) -> int:
            # A 1-bit dot product is bind-then-count: agree - disagree.
            return d - 2 * (a ^ b).bit_count()

        def mem_stat() -> int:
            return int(store.stat()["cells_used"])

        # HDC-native arithmetic, exposed to agents: a ripple-carry circuit over
        # bit atoms rather than a decode-and-add.
        from .adder import HyperAdder

        adder = HyperAdder(d=d)

        def hdc_add(a: int, b: int) -> int:
            return adder.add_hv(a, b)

        def hdc_sub(a: int, b: int) -> int:
            return adder.sub_hv(a, b)

        def agent_echo(n: int) -> int:
            return int(n)

        self.register("hdc.similarity", "hdc/bind/similarity", (("a", "hv"), ("b", "hv")),
                      "Similarity of two hypervectors, scaled to milli-units.", hdc_similarity)
        self.register("hdc.bind", "hdc/bind/construct", (("a", "hv"), ("b", "hv")),
                      "XOR-bind two hypervectors.", hdc_bind, returns="hv")
        self.register("hdc.unbind", "hdc/bind/invert", (("a", "hv"), ("b", "hv")),
                      "Unbind: XOR is self-inverse, so this is bind.", hdc_bind, returns="hv")
        self.register("hdc.bundle", "hdc/bundle/construct", (("a", "hv"), ("b", "hv")),
                      "Majority-bundle two hypervectors.", hdc_bundle, returns="hv")
        self.register("hdc.permute", "hdc/order/rotate", (("v", "hv"), ("k", "int")),
                      "Rotate a hypervector by k bits.", hdc_permute, returns="hv")
        self.register("hdc.cleanup", "hdc/memory/cleanup", (("v", "hv"),),
                      "Index of the stored symbol nearest to a hypervector.", hdc_cleanup)
        self.register("hdc.recall", "hdc/memory/search", (("v", "hv"), ("k", "int")),
                      "Index of the stored cell most similar to a hypervector.", hdc_recall)
        self.register("hdc.encode_int", "hdc/codec/encode", (("n", "int"),),
                      "Encode an integer as a hypervector.", hdc_encode_int, returns="hv")
        self.register("hdc.decode_int", "hdc/codec/decode", (("v", "hv"),),
                      "Decode an integer hypervector.", hdc_decode_int)
        self.register("bit.popcount", "bit/logic/count", (("v", "hv"),),
                      "Population count of a hypervector.", bit_popcount)
        self.register("bit.hamming", "bit/logic/distance", (("a", "hv"), ("b", "hv")),
                      "Hamming distance between two hypervectors.", bit_hamming)
        self.register("net.dot1bit", "bitnet/inference/dot", (("a", "hv"), ("b", "hv")),
                      "1-bit dot product as agreement minus disagreement.", net_dot1bit)
        self.register("mem.stat", "storage/memory/stats", (),
                      "Number of occupied HDC cells.", mem_stat)
        self.register("hdc.add", "hdc/arithmetic/add", (("a", "hv"), ("b", "hv")),
                      "Add two integer hypervectors with a ripple-carry HDC circuit.",
                      hdc_add, returns="hv")
        self.register("hdc.sub", "hdc/arithmetic/subtract", (("a", "hv"), ("b", "hv")),
                      "Subtract two integer hypervectors (two's complement, HDC circuit).",
                      hdc_sub, returns="hv")
        self.register("agent.echo", "agent/plumbing/echo", (("n", "int"),),
                      "Return the argument unchanged; useful for wiring tests.", agent_echo)


class ToolError(Exception):
    """A tool could not be resolved or invoked."""