"""Semantic multi-role relation schemas, not geometric simplices.

A Lean declaration is a bundle of roles: the objects it is about (``(G : Type*)``),
the facts it assumes (``(h : a ∣ b)``), and the goal it asserts or the thing it
defines. A ``Simplex`` is such a bundle with at least two roles. The same schema can
be read off the prose of a book, so the two can be compared role by role.

Design notes (from the archived ``schema-simplex`` prototype):

* a role may be *constrained* to a closed set of allowed values (``choices``);
* a simplex is itself a role, so explicitly supplied schemas can be composed;
  nesting does not discover a proof dependency;
* validation lives in ``__init__``. The prototype put it in a dataclass
  ``__post_init__`` next to a hand-written ``__init__``, where it never ran.
"""
from __future__ import annotations

from collections.abc import Sequence
from typing import Any

KIND_OBJECT = "object"  # something the statement is about: a group G, an integer n
KIND_FACT = "fact"      # something assumed: n > 0, G is finite
KIND_GOAL = "goal"      # what is asserted, or asked to be shown
KIND_OWN = "own"        # what a definition is
KIND_SIMPLEX = "simplex"
KINDS = (KIND_OBJECT, KIND_FACT, KIND_GOAL, KIND_OWN, KIND_SIMPLEX)

MIN_ROLES = 2


class Role:
    """A named slot with a kind, a description and optionally a closed value set."""

    def __init__(
        self,
        name: str,
        kind: str = KIND_OBJECT,
        *,
        type_text: str = "",
        context: str = "",
        choices: Sequence[Any] | None = None,
        attrs: dict[str, Any] | None = None,
    ) -> None:
        if not isinstance(name, str) or not name.strip():
            raise ValueError("a role needs a name")
        if kind not in KINDS:
            raise ValueError(f"unknown role kind {kind!r}; expected one of {KINDS}")
        self.name = name
        self.kind = kind
        self.type_text = type_text
        self.context = context
        if isinstance(choices, (str, bytes)):
            raise ValueError("choices must be a value sequence, not a string")
        self.choices: list[Any] | None = list(choices) if choices is not None else None
        self.attrs: dict[str, Any] = dict(attrs or {})

    def is_constrained(self) -> bool:
        return self.choices is not None

    def validate_value(self, value: Any) -> bool:
        return self.choices is None or value in self.choices

    def check_value(self, value: Any) -> None:
        if not self.validate_value(value):
            raise ValueError(
                f"{value!r} is not allowed for role {self.name!r}; allowed: {self.choices}"
            )

    def as_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {"name": self.name, "kind": self.kind}
        if self.type_text:
            out["type"] = self.type_text
        if self.context:
            out["context"] = self.context
        if self.choices is not None:
            out["choices"] = self.choices
        if self.attrs:
            out["attrs"] = self.attrs
        return out

    def __repr__(self) -> str:
        return f"<Role {self.kind}:{self.name}>"


class Simplex(Role):
    """A reusable bundle of at least two roles, keyed by alias. Itself a role."""

    def __init__(self, name: str, roles: dict[str, Role], *, context: str = "",
                 attrs: dict[str, Any] | None = None) -> None:
        super().__init__(name, KIND_SIMPLEX, context=context, attrs=attrs)
        if len(roles) < MIN_ROLES:
            raise ValueError(
                f"simplex {name!r} needs at least {MIN_ROLES} roles, got {len(roles)}"
            )
        if any(not isinstance(alias, str) or not alias or not isinstance(role, Role)
               for alias, role in roles.items()):
            raise ValueError("simplex slots need nonempty aliases and Role values")
        for role in roles.values():
            references = role.attrs.get("references", [])
            if (not isinstance(references, list)
                    or any(not isinstance(alias, str) or alias not in roles for alias in references)):
                raise ValueError("role references must name slots in this simplex")
        self.roles = dict(roles)

    def by_kind(self, kind: str) -> list[Role]:
        return [role for role in self.roles.values() if role.kind == kind]

    def walk(self) -> list[Role]:
        """Every non-simplex role, descending into nested simplices."""
        out: list[Role] = []
        for role in self.roles.values():
            if isinstance(role, Simplex):
                out.extend(role.walk())
            else:
                out.append(role)
        return out

    def as_dict(self) -> dict[str, Any]:
        return {**super().as_dict(),
                "roles": {alias: role.as_dict() for alias, role in self.roles.items()}}


class State:
    """A role together with a value; unresolved until it has one."""

    def __init__(self, role: Role, value: Any = None) -> None:
        if value is not None:
            role.check_value(value)
        self.role = role
        self.value = value

    def is_resolved(self) -> bool:
        return self.value is not None


class Fact(State):
    """Known at the start."""

    def __init__(self, role: Role, value: Any) -> None:
        if value is None:
            raise ValueError("a Fact must carry a value")
        super().__init__(role, value)


class Goal(State):
    """To be resolved."""


def role_from_dict(data: dict[str, Any]) -> Role:
    """Restore JSON role/schema data without executing code or loading pickle."""
    kind = data.get("kind", KIND_SIMPLEX if "roles" in data else KIND_OBJECT)
    if kind == KIND_SIMPLEX:
        return Simplex(data["name"], {alias: role_from_dict(role)
                                   for alias, role in data["roles"].items()},
                       context=data.get("context", ""), attrs=data.get("attrs"))
    return Role(data["name"], kind, type_text=data.get("type", ""),
                context=data.get("context", ""), choices=data.get("choices"),
                attrs=data.get("attrs"))


class SimplexInstance:
    """Values bind local slot aliases; a resolved goal is supplied, not proved here."""

    def __init__(self, schema: Simplex, facts: dict[str, Any] | None = None,
                 goals: Sequence[str] | None = None) -> None:
        self.schema = schema
        facts = facts or {}
        goal_aliases = list(goals) if goals is not None else [
            alias for alias, role in schema.roles.items() if role.kind in (KIND_GOAL, KIND_OWN)]
        if (len(set(goal_aliases)) != len(goal_aliases)
                or not (set(facts) | set(goal_aliases)) <= set(schema.roles)
                or set(facts) & set(goal_aliases)):
            raise ValueError("facts/goals need distinct, declared slot aliases")
        self.facts = {alias: Fact(schema.roles[alias], value) for alias, value in facts.items()}
        self.goals = {alias: Goal(schema.roles[alias]) for alias in goal_aliases}

    def resolve(self, alias: str, value: Any) -> None:
        if alias not in self.goals or value is None:
            raise ValueError("resolve needs a declared goal and a non-None value")
        self.goals[alias].role.check_value(value)
        self.goals[alias].value = value

    def is_completed(self) -> bool:
        return bool(self.goals) and all(goal.is_resolved() for goal in self.goals.values())