from app.modules.agent.repository import AgentRepository
from app.modules.agent.engine.router.schemas import RouterContext


class RouterContextStore:
    def __init__(self, repository: AgentRepository) -> None:
        self._repo = repository

    def get(self, conversation_key: str) -> RouterContext:
        return self._repo.get_router_context(conversation_key)

    def update(
        self,
        conversation_key: str,
        *,
        domain_id: str,
        process_id: str,
        user_message: str,
        assistant_message: str,
        decision_type: str = "start",
        max_history: int = 10,
    ) -> None:
        self._repo.update_router_context(
            conversation_key,
            domain_id=domain_id,
            process_id=process_id,
            user_message=user_message,
            assistant_message=assistant_message,
            decision_type=decision_type,
            max_history=max_history,
        )
