from __future__ import annotations

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


StateT = TypeVar("StateT")
CandidateT = TypeVar("CandidateT")
MetadataT = TypeVar("MetadataT")


@dataclass(frozen=True)
class StepTrace(Generic[StateT, CandidateT, MetadataT]):
    previous_state: StateT
    next_state: StateT
    candidate: CandidateT
    accepted: bool
    elapsed_seconds: float
    metadata: MetadataT | None = None


@dataclass(frozen=True)
class ExecutionRecord(Generic[StateT, CandidateT, MetadataT]):
    accepted: tuple[StepTrace[StateT, CandidateT, MetadataT], ...]
    attempts: tuple[StepTrace[StateT, CandidateT, MetadataT], ...]
    final_state: StateT
    total_seconds: float

    @property
    def accepted_count(self) -> int:
        return len(self.accepted)

    @property
    def attempt_count(self) -> int:
        return len(self.attempts)


class AcceptRejectSystem(Protocol[StateT, CandidateT, MetadataT]):
    def step(self, state: StateT) -> StepTrace[StateT, CandidateT, MetadataT]:
        ...


def run_until_acceptance(
    system: AcceptRejectSystem[StateT, CandidateT, MetadataT],
    initial_state: StateT,
    *,
    max_attempts: int,
) -> ExecutionRecord[StateT, CandidateT, MetadataT]:
    start_time = time.perf_counter()
    attempts: list[StepTrace[StateT, CandidateT, MetadataT]] = []
    current_state = initial_state
    for _ in range(max_attempts):
        step = system.step(current_state)
        attempts.append(step)
        if step.accepted:
            return ExecutionRecord(
                accepted=(step,),
                attempts=tuple(attempts),
                final_state=step.next_state,
                total_seconds=time.perf_counter() - start_time,
            )
        current_state = step.next_state
    raise RuntimeError("failed to produce an accepted step within max_attempts")


def run_until_acceptance_count(
    system: AcceptRejectSystem[StateT, CandidateT, MetadataT],
    initial_state: StateT,
    *,
    accepted_count: int,
    max_attempts_per_accept: int,
) -> ExecutionRecord[StateT, CandidateT, MetadataT]:
    start_time = time.perf_counter()
    attempts: list[StepTrace[StateT, CandidateT, MetadataT]] = []
    accepted_steps: list[StepTrace[StateT, CandidateT, MetadataT]] = []
    current_state = initial_state
    for _ in range(accepted_count):
        accepted_run = run_until_acceptance(
            system,
            current_state,
            max_attempts=max_attempts_per_accept,
        )
        attempts.extend(accepted_run.attempts)
        accepted_steps.extend(accepted_run.accepted)
        current_state = accepted_run.final_state
    return ExecutionRecord(
        accepted=tuple(accepted_steps),
        attempts=tuple(attempts),
        final_state=current_state,
        total_seconds=time.perf_counter() - start_time,
    )


def merge_execution_records(
    left: ExecutionRecord[StateT, CandidateT, MetadataT],
    right: ExecutionRecord[StateT, CandidateT, MetadataT],
) -> ExecutionRecord[StateT, CandidateT, MetadataT]:
    if left.final_state != right.attempts[0].previous_state:
        raise ValueError("right record does not continue from left final_state")
    return ExecutionRecord(
        accepted=left.accepted + right.accepted,
        attempts=left.attempts + right.attempts,
        final_state=right.final_state,
        total_seconds=left.total_seconds + right.total_seconds,
    )
