"""Composition algebra for π*.

Per ticket #000015 §2.3, a composition ``π*_a ∘ π*_b`` (apply π*_b
first, then π*_a) is a valid canonical projection iff:

1. **Type-compatible** — ``π*_b.codomain`` ⊆ ``π*_a.domain``.
2. **Determinism preserved** — both factors are deterministic ⇒ chain
   is deterministic.
3. **Equivalence-class-preserving** — function composition trivially
   preserves classes.

Each composition is itself a registered π* with key
``"<inner>-then-<outer>@<version>"`` so callers see compositions as
first-class projections, not implementation details.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from typing import Optional

from arborist.pi_star.protocol import PiStar, PiStarError, registry_key
from arborist.pi_star.registry import REGISTRY, register


@dataclass(frozen=True)
class PiStarComposition:
    """A composition ``outer ∘ inner`` exposed as a π*.

    ``canonicalize(raw)`` runs ``inner.canonicalize`` first, then
    ``outer.canonicalize`` on the result. Domain is the inner π*'s
    domain (we accept what the first stage accepts); codomain is the
    outer π*'s codomain.
    """

    name: str
    version: str
    inner_key: str
    outer_key: str
    inner: PiStar
    outer: PiStar

    @property
    def domain(self) -> str:
        return self.inner.domain

    def canonicalize(self, raw: bytes) -> bytes:
        try:
            mid = self.inner.canonicalize(raw)
        except PiStarError:
            raise
        return self.outer.canonicalize(mid)


def canonical_composition_id(inner_key: str, outer_key: str) -> str:
    """Stable SHA-256 fingerprint for a composition manifest.

    The composition's identity is the canonical-JSON of its (inner,
    outer) keys. Useful for governance hash inclusion when a policy
    pins which composition is in effect.
    """
    payload = json.dumps(
        {"chain": [inner_key, outer_key]},
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def compose(
    inner_key: str,
    outer_key: str,
    *,
    name: Optional[str] = None,
    version: str = "v1",
    register_in_registry: bool = True,
) -> PiStarComposition:
    """Build (and optionally register) a composition.

    Default name is ``"<inner-name>-then-<outer-name>"`` so the
    canonical example renders as ``"wikitext-base-then-claim-lattice"``.
    """
    if inner_key not in REGISTRY:
        raise KeyError(f"unknown inner π*: {inner_key}")
    if outer_key not in REGISTRY:
        raise KeyError(f"unknown outer π*: {outer_key}")
    inner = REGISTRY[inner_key]
    outer = REGISTRY[outer_key]
    derived_name = name or f"{inner.name}-then-{outer.name}"

    composition = PiStarComposition(
        name=derived_name,
        version=version,
        inner_key=inner_key,
        outer_key=outer_key,
        inner=inner,
        outer=outer,
    )
    if register_in_registry:
        register(composition)
    return composition
