from __future__ import annotations

from dataclasses import dataclass
import time
from typing import Generic, TypeVar

from .architectures import (
    AcceptancePolicy,
    Categorizer,
    Critic,
    Generator,
    PolicyDecision,
    StateTransition,
)
from .runtime import StepTrace


StateT = TypeVar("StateT")
CandidateT = TypeVar("CandidateT")
CritiqueT = TypeVar("CritiqueT")
CategoryT = TypeVar("CategoryT")


@dataclass(frozen=True)
class HybridStepMetadata(Generic[CritiqueT, CategoryT]):
    critique: CritiqueT
    category: CategoryT
    decision: PolicyDecision


class CooperativeSystem(Generic[StateT, CandidateT, CritiqueT, CategoryT]):
    def __init__(
        self,
        *,
        generator: Generator[StateT, CandidateT],
        critic: Critic[StateT, CandidateT, CritiqueT],
        categorizer: Categorizer[StateT, CandidateT, CategoryT],
        policy: AcceptancePolicy[StateT, CandidateT, CritiqueT, CategoryT],
        transition: StateTransition[StateT, CandidateT, CritiqueT, CategoryT],
    ) -> None:
        self.generator = generator
        self.critic = critic
        self.categorizer = categorizer
        self.policy = policy
        self.transition = transition

    def step(
        self,
        state: StateT,
    ) -> StepTrace[StateT, CandidateT, HybridStepMetadata[CritiqueT, CategoryT]]:
        start_time = time.perf_counter()
        candidate = self.generator.generate(state)
        critique = self.critic.critique(state, candidate)
        category = self.categorizer.categorize(state, candidate)
        decision = self.policy.decide(state, candidate, critique, category)
        next_state = self.transition.advance(state, candidate, critique, category, decision)
        return StepTrace(
            previous_state=state,
            next_state=next_state,
            candidate=candidate,
            accepted=decision.accepted,
            elapsed_seconds=time.perf_counter() - start_time,
            metadata=HybridStepMetadata(
                critique=critique,
                category=category,
                decision=decision,
            ),
        )
