"""Small linearizability checker for branch-head histories (stdlib only).

Method: Wing & Gong search with Lowe-style memoisation of (linearized-set,

model-state), the approach Porcupine (MIT, github.com/anishathalye/porcupine)

implements; this is an independent re-implementation, no Porcupine code is

copied. Histories are partitioned per branch (P-compositionality): branch

heads are independent registers.

Model (fast-forward-only branch head, what `oak push` / `git push` without

--force promise):

- ``push(c) -> ok``      legal iff head is None or head is an ancestor of c

                         (c == head included); head := c.

- ``push(c) -> rejected`` no effect (a rejected push must not move the head).

- ``push(c) -> unknown``  the client died / timed out / lost the reply: it may

                         take effect at any point after its call, or never.

- ``read -> h``           legal iff head == h (an observed server head).

Ancestry comes from the union of independently-read commit graphs.

"""

from __future__ import annotations

import math

from dataclasses import dataclass, field

from typing import Callable, Hashable, Iterable, Optional

OK = "ok"

REJECTED = "rejected"

UNKNOWN = "unknown"

@dataclass

class Op:

    client: str

    kind: str  # "push" | "read"

    value: Optional[str]  # pushed commit, or observed head

    call: float

    ret: float  # math.inf for unknown outcomes

    outcome: str = OK

    branch: str = "main"

    op_id: int = -1

    meta: dict = field(default_factory=dict)

@dataclass

class CheckResult:

    linearizable: Optional[bool]  # None = search budget exhausted (unmeasured)

    order: list[int]

    explored: int

    reason: str = ""

def _next_states(state: Optional[str], op: Op, is_ancestor: Callable[[str, str], bool]) -> list[Optional[str]]:

    if op.kind == "read":

        return [state] if state == op.value else []

    if op.kind != "push":

        raise ValueError(f"unknown op kind {op.kind!r}")

    ff = op.value is not None and (state is None or state == op.value or is_ancestor(state, op.value))

    if op.outcome == OK:

        return [op.value] if ff else []

    if op.outcome == REJECTED:

        return [state]

    # unknown: either it never took effect, or it did (only legal as a fast-forward)

    return [state, op.value] if ff and op.value != state else [state]

def check(ops: list[Op], is_ancestor: Callable[[str, str], bool], initial: Optional[str] = None,

          budget: int = 2_000_000) -> CheckResult:

    """Is there a total order of ``ops`` respecting real time and the model?"""

    ops = list(ops)

    for index, op in enumerate(ops):

        op.op_id = index

    n = len(ops)

    must = 0

    for op in ops:

        if op.ret != math.inf:

            must |= 1 << op.op_id

    explored = 0

    seen: set[tuple[int, Hashable]] = set()

    order: list[int] = []

    def candidates(mask: int) -> list[Op]:

        pending = [op for op in ops if not (mask >> op.op_id) & 1]

        if not pending:

            return []

        horizon = min(op.ret for op in pending)

        return [op for op in pending if op.call <= horizon]

    # Iterative DFS to avoid recursion limits.

    stack: list[tuple[int, Optional[str], list, int]] = []  # (mask, state, choices, choice_index)

    def expand(mask: int, state: Optional[str]) -> list[tuple[int, Optional[str]]]:

        result = []

        for op in candidates(mask):

            for nxt in _next_states(state, op, is_ancestor):

                result.append((op.op_id, nxt))

        return result

    mask, state = 0, initial

    if must == 0:

        return CheckResult(True, [], 0)

    stack.append((mask, state, expand(mask, state), 0))

    seen.add((mask, state))

    while stack:

        mask, state, choices, idx = stack[-1]

        if (mask & must) == must:

            return CheckResult(True, list(order), explored)

        if idx >= len(choices):

            stack.pop()

            if order:

                order.pop()

            continue

        stack[-1] = (mask, state, choices, idx + 1)

        op_id, nxt = choices[idx]

        new_mask = mask | (1 << op_id)

        key = (new_mask, nxt)

        explored += 1

        if explored > budget:

            return CheckResult(None, list(order), explored, "search budget exhausted")

        if key in seen:

            continue

        seen.add(key)

        order.append(op_id)

        stack.append((new_mask, nxt, expand(new_mask, nxt), 0))

    return CheckResult(False, [], explored, "no linearization respects real time and the fast-forward model")

def check_partitioned(ops: Iterable[Op], is_ancestor: Callable[[str, str], bool],

                      initial: Optional[dict[str, Optional[str]]] = None) -> dict[str, CheckResult]:

    by_branch: dict[str, list[Op]] = {}

    for op in ops:

        by_branch.setdefault(op.branch, []).append(op)

    initial = initial or {}

    return {branch: check(branch_ops, is_ancestor, initial.get(branch)) for branch, branch_ops in by_branch.items()}

def lost_acknowledged(ops: Iterable[Op], final_heads: dict[str, Optional[str]],

                      is_ancestor: Callable[[str, str], bool]) -> list[Op]:

    """Acknowledged pushes whose commit is not contained in the branch's final head."""

    lost = []

    for op in ops:

        if op.kind != "push" or op.outcome != OK or op.value is None:

            continue

        final = final_heads.get(op.branch)

        if final is None or not (final == op.value or is_ancestor(op.value, final)):

            lost.append(op)

    return lost

def ancestry_from_parents(parents: dict[str, tuple]) -> Callable[[str, str], bool]:

    cache: dict[str, frozenset] = {}

    def ancestors(commit: str) -> frozenset:

        if commit in cache:

            return cache[commit]

        seen = set()

        stack = [commit]

        while stack:

            current = stack.pop()

            if current in seen:

                continue

            seen.add(current)

            stack.extend(p for p in parents.get(current, ()) if p)

        cache[commit] = frozenset(seen)

        return cache[commit]

    def is_ancestor(ancestor: str, descendant: str) -> bool:

        return ancestor in ancestors(descendant)

    return is_ancestor