from __future__ import annotations

from epochler.contract.models import Conversation, ConversationState


_TRANSITIONS: dict[ConversationState, set[ConversationState]] = {
    ConversationState.DRAFT: {
        ConversationState.NEEDS_INPUT,
        ConversationState.READY_TO_PREPROCESS,
        ConversationState.READY_TO_TRAIN,
        ConversationState.RUNNING,
        ConversationState.STOPPED,
    },
    ConversationState.NEEDS_INPUT: {
        ConversationState.DRAFT,
        ConversationState.READY_TO_PREPROCESS,
        ConversationState.READY_TO_TRAIN,
        ConversationState.STOPPED,
    },
    ConversationState.READY_TO_PREPROCESS: {
        ConversationState.RUNNING,
        ConversationState.NEEDS_INPUT,
        ConversationState.STOPPED,
    },
    ConversationState.READY_TO_TRAIN: {
        ConversationState.RUNNING,
        ConversationState.NEEDS_INPUT,
        ConversationState.STOPPED,
    },
    ConversationState.RUNNING: {
        ConversationState.PAUSED,
        ConversationState.REVIEWING_RESULTS,
        ConversationState.FAILED,
        ConversationState.STOPPED,
    },
    ConversationState.PAUSED: {
        ConversationState.RUNNING,
        ConversationState.STOPPED,
    },
    ConversationState.REVIEWING_RESULTS: {
        ConversationState.READY_TO_TRAIN,
        ConversationState.COMPLETED,
        ConversationState.STOPPED,
    },
    ConversationState.FAILED: {
        ConversationState.READY_TO_PREPROCESS,
        ConversationState.READY_TO_TRAIN,
        ConversationState.STOPPED,
    },
    ConversationState.COMPLETED: set(),
    ConversationState.STOPPED: set(),
}


def can_transition(current: ConversationState, target: ConversationState) -> bool:
    return target in _TRANSITIONS.get(current, set())


def transition(conv: Conversation, target: ConversationState) -> Conversation:
    if not can_transition(conv.state, target):
        raise ValueError(f"Invalid state transition: {conv.state.value} -> {target.value}")
    conv.state = target
    return conv


def infer_state_from_contract(conv: Conversation) -> ConversationState:
    return conv.readiness.current_stage
