from __future__ import annotations

from dataclasses import dataclass

from synaptopus.architectures import PolicyDecision
from synaptopus.orchestration import CooperativeSystem
from synaptopus.runtime import run_until_acceptance, run_until_acceptance_count


@dataclass(frozen=True)
class SequenceState:
    accepted: tuple[int, ...] = ()
    attempts: int = 0


class IncrementingGenerator:
    def generate(self, state: SequenceState) -> int:
        return state.attempts + 1


class EvenCritic:
    def critique(self, state: SequenceState, candidate: int) -> bool:
        return candidate % 2 == 0


class ModuloCategorizer:
    def categorize(self, state: SequenceState, candidate: int) -> str:
        return "novel" if candidate % 3 else "repeat"


class AcceptEvenNovelPolicy:
    def decide(
        self,
        state: SequenceState,
        candidate: int,
        critique: bool,
        category: str,
    ) -> PolicyDecision:
        accepted = critique and category == "novel"
        label = "accept" if accepted else "reject"
        return PolicyDecision(accepted=accepted, label=label)


class SequenceTransition:
    def advance(
        self,
        state: SequenceState,
        candidate: int,
        critique: bool,
        category: str,
        decision: PolicyDecision,
    ) -> SequenceState:
        if decision.accepted:
            return SequenceState(
                accepted=state.accepted + (candidate,),
                attempts=state.attempts + 1,
            )
        return SequenceState(
            accepted=state.accepted,
            attempts=state.attempts + 1,
        )


def build_system() -> CooperativeSystem[SequenceState, int, bool, str]:
    return CooperativeSystem(
        generator=IncrementingGenerator(),
        critic=EvenCritic(),
        categorizer=ModuloCategorizer(),
        policy=AcceptEvenNovelPolicy(),
        transition=SequenceTransition(),
    )


def test_cooperative_system_exposes_component_metadata() -> None:
    step = build_system().step(SequenceState())

    assert step.candidate == 1
    assert step.accepted is False
    assert step.metadata is not None
    assert step.metadata.critique is False
    assert step.metadata.category == "novel"
    assert step.metadata.decision.label == "reject"


def test_cooperative_system_runs_until_acceptance() -> None:
    record = run_until_acceptance(build_system(), SequenceState(), max_attempts=5)

    assert record.accepted_count == 1
    assert record.attempt_count == 2
    assert record.final_state.accepted == (2,)
    assert record.accepted[0].metadata is not None
    assert record.accepted[0].metadata.decision.label == "accept"


def test_cooperative_system_runs_multiple_acceptances() -> None:
    record = run_until_acceptance_count(
        build_system(),
        SequenceState(),
        accepted_count=3,
        max_attempts_per_accept=6,
    )

    assert record.final_state.accepted == (2, 4, 8)
    assert tuple(step.candidate for step in record.accepted) == (2, 4, 8)