ss-tools/backend/src/services/dataset_review/clarification_engine.py

# [DEF:ClarificationEngine:Module]
# @COMPLEXITY: 4
# @SEMANTICS: dataset_review, clarification, question_payload, answer_persistence, readiness, findings
# @PURPOSE: Manage one-question-at-a-time clarification state, deterministic answer persistence, and readiness/finding updates.
# @LAYER: Domain
# @RELATION: [DEPENDS_ON] ->[DatasetReviewSessionRepository]
# @RELATION: [DEPENDS_ON] ->[ClarificationSession]
# @RELATION: [DEPENDS_ON] ->[ClarificationQuestion]
# @RELATION: [DEPENDS_ON] ->[ClarificationAnswer]
# @RELATION: [DEPENDS_ON] ->[ValidationFinding]
# @PRE: Target session contains a persisted clarification aggregate in the current ownership scope.
# @POST: Active clarification payload exposes one highest-priority unresolved question, and each recorded answer is persisted before pointer/readiness mutation.
# @SIDE_EFFECT: Persists clarification answers, question/session states, and related readiness/finding changes.
# @DATA_CONTRACT: Input[DatasetReviewSession|ClarificationAnswerCommand] -> Output[ClarificationStateResult]
# @INVARIANT: Only one active clarification question may exist at a time; skipped and expert-review items remain unresolved and visible.
from __future__ import annotations

# [DEF:imports:Block]
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import List, Optional

from src.core.logger import belief_scope, logger
from src.models.auth import User
from src.models.dataset_review import (
    AnswerKind,
    ClarificationAnswer,
    ClarificationQuestion,
    ClarificationSession,
    ClarificationStatus,
    DatasetReviewSession,
    FindingArea,
    FindingSeverity,
    QuestionState,
    ReadinessState,
    RecommendedAction,
    ResolutionState,
    SessionPhase,
    ValidationFinding,
)
from src.services.dataset_review.repositories.session_repository import (
    DatasetReviewSessionRepository,
)
# [/DEF:imports:Block]


# [DEF:ClarificationQuestionPayload:Class]
# @COMPLEXITY: 2
# @PURPOSE: Typed active-question payload returned to the API layer.
@dataclass
class ClarificationQuestionPayload:
    question_id: str
    clarification_session_id: str
    topic_ref: str
    question_text: str
    why_it_matters: str
    current_guess: Optional[str]
    priority: int
    state: QuestionState
    options: list[dict[str, object]] = field(default_factory=list)


# [/DEF:ClarificationQuestionPayload:Class]


# [DEF:ClarificationStateResult:Class]
# @COMPLEXITY: 2
# @PURPOSE: Clarification state result carrying the current session, active payload, and changed findings.
@dataclass
class ClarificationStateResult:
    clarification_session: ClarificationSession
    current_question: Optional[ClarificationQuestionPayload]
    session: DatasetReviewSession
    changed_findings: List[ValidationFinding] = field(default_factory=list)


# [/DEF:ClarificationStateResult:Class]


# [DEF:ClarificationAnswerCommand:Class]
# @COMPLEXITY: 2
# @PURPOSE: Typed answer command for clarification state mutation.
@dataclass
class ClarificationAnswerCommand:
    session: DatasetReviewSession
    question_id: str
    answer_kind: AnswerKind
    answer_value: Optional[str]
    user: User


# [/DEF:ClarificationAnswerCommand:Class]


# [DEF:ClarificationEngine:Class]
# @COMPLEXITY: 4
# @PURPOSE: Provide deterministic one-question-at-a-time clarification selection and answer persistence.
# @RELATION: [DEPENDS_ON] ->[DatasetReviewSessionRepository]
# @RELATION: [DEPENDS_ON] ->[ClarificationSession]
# @RELATION: [DEPENDS_ON] ->[ValidationFinding]
# @PRE: Repository is bound to the current request transaction scope.
# @POST: Returned clarification state is persistence-backed and aligned with session readiness/recommended action.
# @SIDE_EFFECT: Mutates clarification answers, session flags, and related clarification findings.
class ClarificationEngine:
    # [DEF:ClarificationEngine_init:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Bind repository dependency for clarification persistence operations.
    def __init__(self, repository: DatasetReviewSessionRepository) -> None:
        self.repository = repository

    # [/DEF:ClarificationEngine_init:Function]

    # [DEF:build_question_payload:Function]
    # @COMPLEXITY: 4
    # @PURPOSE: Return the one active highest-priority clarification question payload with why-it-matters, current guess, and options.
    # @RELATION: [DEPENDS_ON] ->[ClarificationQuestion]
    # @RELATION: [DEPENDS_ON] ->[ClarificationOption]
    # @PRE: Session contains unresolved clarification state or a resumable clarification session.
    # @POST: Returns exactly one active/open question payload or None when no unresolved question remains.
    # @SIDE_EFFECT: Normalizes the active-question pointer and clarification status in persistence.
    # @DATA_CONTRACT: Input[DatasetReviewSession] -> Output[ClarificationQuestionPayload|None]
    def build_question_payload(
        self,
        session: DatasetReviewSession,
    ) -> Optional[ClarificationQuestionPayload]:
        with belief_scope("ClarificationEngine.build_question_payload"):
            clarification_session = self._get_latest_clarification_session(session)
            if clarification_session is None:
                logger.reason(
                    "Clarification payload requested without clarification session",
                    extra={"session_id": session.session_id},
                )
                return None

            active_questions = [
                question
                for question in clarification_session.questions
                if question.state == QuestionState.OPEN
            ]
            active_questions.sort(
                key=lambda item: (
                    -int(item.priority),
                    item.created_at,
                    item.question_id,
                )
            )

            if not active_questions:
                clarification_session.current_question_id = None
                clarification_session.status = ClarificationStatus.COMPLETED
                session.readiness_state = self._derive_readiness_state(session)
                session.recommended_action = self._derive_recommended_action(session)
                if session.current_phase == SessionPhase.CLARIFICATION:
                    session.current_phase = SessionPhase.REVIEW
                self.repository.db.commit()
                logger.reflect(
                    "No unresolved clarification question remains",
                    extra={"session_id": session.session_id},
                )
                return None

            selected_question = active_questions[0]
            clarification_session.current_question_id = selected_question.question_id
            clarification_session.status = ClarificationStatus.ACTIVE
            session.readiness_state = ReadinessState.CLARIFICATION_ACTIVE
            session.recommended_action = RecommendedAction.ANSWER_NEXT_QUESTION
            session.current_phase = SessionPhase.CLARIFICATION

            logger.reason(
                "Selected active clarification question",
                extra={
                    "session_id": session.session_id,
                    "clarification_session_id": clarification_session.clarification_session_id,
                    "question_id": selected_question.question_id,
                    "priority": selected_question.priority,
                },
            )
            self.repository.db.commit()

            payload = ClarificationQuestionPayload(
                question_id=selected_question.question_id,
                clarification_session_id=selected_question.clarification_session_id,
                topic_ref=selected_question.topic_ref,
                question_text=selected_question.question_text,
                why_it_matters=selected_question.why_it_matters,
                current_guess=selected_question.current_guess,
                priority=selected_question.priority,
                state=selected_question.state,
                options=[
                    {
                        "option_id": option.option_id,
                        "question_id": option.question_id,
                        "label": option.label,
                        "value": option.value,
                        "is_recommended": option.is_recommended,
                        "display_order": option.display_order,
                    }
                    for option in sorted(
                        selected_question.options,
                        key=lambda item: (
                            item.display_order,
                            item.label,
                            item.option_id,
                        ),
                    )
                ],
            )
            logger.reflect(
                "Clarification payload built",
                extra={
                    "session_id": session.session_id,
                    "question_id": payload.question_id,
                    "option_count": len(payload.options),
                },
            )
            return payload

    # [/DEF:build_question_payload:Function]

    # [DEF:record_answer:Function]
    # @COMPLEXITY: 4
    # @PURPOSE: Persist one clarification answer before any pointer/readiness mutation and compute deterministic state impact.
    # @RELATION: [DEPENDS_ON] ->[ClarificationAnswer]
    # @RELATION: [DEPENDS_ON] ->[ValidationFinding]
    # @PRE: Target question belongs to the session's active clarification session and is still open.
    # @POST: Answer row is persisted before current-question pointer advances; skipped/expert-review items remain unresolved and visible.
    # @SIDE_EFFECT: Inserts answer row, mutates question/session states, updates clarification findings, and commits.
    # @DATA_CONTRACT: Input[ClarificationAnswerCommand] -> Output[ClarificationStateResult]
    def record_answer(
        self, command: ClarificationAnswerCommand
    ) -> ClarificationStateResult:
        with belief_scope("ClarificationEngine.record_answer"):
            session = command.session
            clarification_session = self._get_latest_clarification_session(session)
            if clarification_session is None:
                logger.explore(
                    "Cannot record clarification answer because no clarification session exists",
                    extra={"session_id": session.session_id},
                )
                raise ValueError("Clarification session not found")

            question = self._find_question(clarification_session, command.question_id)
            if question is None:
                logger.explore(
                    "Cannot record clarification answer for foreign or missing question",
                    extra={
                        "session_id": session.session_id,
                        "question_id": command.question_id,
                    },
                )
                raise ValueError("Clarification question not found")

            if question.answer is not None:
                logger.explore(
                    "Rejected duplicate clarification answer submission",
                    extra={
                        "session_id": session.session_id,
                        "question_id": command.question_id,
                    },
                )
                raise ValueError("Clarification question already answered")

            if (
                clarification_session.current_question_id
                and clarification_session.current_question_id != question.question_id
            ):
                logger.explore(
                    "Rejected answer for non-active clarification question",
                    extra={
                        "session_id": session.session_id,
                        "question_id": question.question_id,
                        "current_question_id": clarification_session.current_question_id,
                    },
                )
                raise ValueError(
                    "Only the active clarification question can be answered"
                )

            normalized_answer_value = self._normalize_answer_value(
                command.answer_kind, command.answer_value, question
            )

            logger.reason(
                "Persisting clarification answer before state advancement",
                extra={
                    "session_id": session.session_id,
                    "question_id": question.question_id,
                    "answer_kind": command.answer_kind.value,
                },
            )
            persisted_answer = ClarificationAnswer(
                question_id=question.question_id,
                answer_kind=command.answer_kind,
                answer_value=normalized_answer_value,
                answered_by_user_id=command.user.id,
                impact_summary=self._build_impact_summary(
                    question, command.answer_kind, normalized_answer_value
                ),
            )
            self.repository.db.add(persisted_answer)
            self.repository.db.flush()

            changed_finding = self._upsert_clarification_finding(
                session=session,
                question=question,
                answer_kind=command.answer_kind,
                answer_value=normalized_answer_value,
            )

            if command.answer_kind == AnswerKind.SELECTED:
                question.state = QuestionState.ANSWERED
            elif command.answer_kind == AnswerKind.CUSTOM:
                question.state = QuestionState.ANSWERED
            elif command.answer_kind == AnswerKind.SKIPPED:
                question.state = QuestionState.SKIPPED
            elif command.answer_kind == AnswerKind.EXPERT_REVIEW:
                question.state = QuestionState.EXPERT_REVIEW

            question.updated_at = datetime.utcnow()
            self.repository.db.flush()

            clarification_session.resolved_count = self._count_resolved_questions(
                clarification_session
            )
            clarification_session.remaining_count = self._count_remaining_questions(
                clarification_session
            )
            clarification_session.summary_delta = self.summarize_progress(
                clarification_session
            )
            clarification_session.updated_at = datetime.utcnow()

            next_question = self._select_next_open_question(clarification_session)
            clarification_session.current_question_id = (
                next_question.question_id if next_question else None
            )
            clarification_session.status = (
                ClarificationStatus.ACTIVE
                if next_question
                else ClarificationStatus.COMPLETED
            )
            if clarification_session.status == ClarificationStatus.COMPLETED:
                clarification_session.completed_at = datetime.utcnow()

            session.readiness_state = self._derive_readiness_state(session)
            session.recommended_action = self._derive_recommended_action(session)
            session.current_phase = (
                SessionPhase.CLARIFICATION
                if clarification_session.current_question_id
                else SessionPhase.REVIEW
            )
            self.repository.bump_session_version(session)

            self.repository.db.commit()
            self.repository.db.refresh(session)

            logger.reflect(
                "Clarification answer recorded and session advanced",
                extra={
                    "session_id": session.session_id,
                    "question_id": question.question_id,
                    "next_question_id": clarification_session.current_question_id,
                    "readiness_state": session.readiness_state.value,
                    "remaining_count": clarification_session.remaining_count,
                },
            )

            return ClarificationStateResult(
                clarification_session=clarification_session,
                current_question=self.build_question_payload(session),
                session=session,
                changed_findings=[changed_finding] if changed_finding else [],
            )

    # [/DEF:record_answer:Function]

    # [DEF:summarize_progress:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Produce a compact progress summary for pause/resume and completion UX.
    # @RELATION: [DEPENDS_ON] ->[ClarificationSession]
    def summarize_progress(self, clarification_session: ClarificationSession) -> str:
        resolved = self._count_resolved_questions(clarification_session)
        remaining = self._count_remaining_questions(clarification_session)
        return f"{resolved} resolved, {remaining} unresolved"

    # [/DEF:summarize_progress:Function]

    # [DEF:_get_latest_clarification_session:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Select the latest clarification session for the current dataset review aggregate.
    def _get_latest_clarification_session(
        self,
        session: DatasetReviewSession,
    ) -> Optional[ClarificationSession]:
        if not session.clarification_sessions:
            return None
        ordered_sessions = sorted(
            session.clarification_sessions,
            key=lambda item: (item.started_at, item.clarification_session_id),
            reverse=True,
        )
        return ordered_sessions[0]

    # [/DEF:_get_latest_clarification_session:Function]

    # [DEF:_find_question:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Resolve a clarification question from the active clarification aggregate.
    def _find_question(
        self,
        clarification_session: ClarificationSession,
        question_id: str,
    ) -> Optional[ClarificationQuestion]:
        for question in clarification_session.questions:
            if question.question_id == question_id:
                return question
        return None

    # [/DEF:_find_question:Function]

    # [DEF:_select_next_open_question:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Select the next unresolved question in deterministic priority order.
    def _select_next_open_question(
        self,
        clarification_session: ClarificationSession,
    ) -> Optional[ClarificationQuestion]:
        open_questions = [
            question
            for question in clarification_session.questions
            if question.state == QuestionState.OPEN
        ]
        if not open_questions:
            return None
        open_questions.sort(
            key=lambda item: (-int(item.priority), item.created_at, item.question_id)
        )
        return open_questions[0]

    # [/DEF:_select_next_open_question:Function]

    # [DEF:_count_resolved_questions:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Count questions whose answers fully resolved the ambiguity.
    def _count_resolved_questions(
        self, clarification_session: ClarificationSession
    ) -> int:
        return sum(
            1
            for question in clarification_session.questions
            if question.state == QuestionState.ANSWERED
        )

    # [/DEF:_count_resolved_questions:Function]

    # [DEF:_count_remaining_questions:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Count questions still unresolved or deferred after clarification interaction.
    def _count_remaining_questions(
        self, clarification_session: ClarificationSession
    ) -> int:
        return sum(
            1
            for question in clarification_session.questions
            if question.state
            in {QuestionState.OPEN, QuestionState.SKIPPED, QuestionState.EXPERT_REVIEW}
        )

    # [/DEF:_count_remaining_questions:Function]

    # [DEF:_normalize_answer_value:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Validate and normalize answer payload based on answer kind and active question options.
    def _normalize_answer_value(
        self,
        answer_kind: AnswerKind,
        answer_value: Optional[str],
        question: ClarificationQuestion,
    ) -> Optional[str]:
        normalized_answer_value = (
            str(answer_value).strip() if answer_value is not None else None
        )
        if (
            answer_kind in {AnswerKind.SELECTED, AnswerKind.CUSTOM}
            and not normalized_answer_value
        ):
            raise ValueError(
                "answer_value is required for selected or custom clarification answers"
            )
        if answer_kind == AnswerKind.SELECTED:
            allowed_values = {option.value for option in question.options}
            if normalized_answer_value not in allowed_values:
                raise ValueError(
                    "answer_value must match one of the current clarification options"
                )
        if answer_kind == AnswerKind.SKIPPED:
            return normalized_answer_value or "skipped"
        if answer_kind == AnswerKind.EXPERT_REVIEW:
            return normalized_answer_value or "expert_review"
        return normalized_answer_value

    # [/DEF:_normalize_answer_value:Function]

    # [DEF:_build_impact_summary:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Build a compact audit note describing how the clarification answer affects session state.
    def _build_impact_summary(
        self,
        question: ClarificationQuestion,
        answer_kind: AnswerKind,
        answer_value: Optional[str],
    ) -> str:
        if answer_kind == AnswerKind.SKIPPED:
            return f"Clarification for {question.topic_ref} was skipped and remains unresolved."
        if answer_kind == AnswerKind.EXPERT_REVIEW:
            return f"Clarification for {question.topic_ref} was deferred for expert review."
        return f"Clarification for {question.topic_ref} recorded as '{answer_value}'."

    # [/DEF:_build_impact_summary:Function]

    # [DEF:_upsert_clarification_finding:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Keep one finding per clarification topic aligned with answer outcome and unresolved visibility rules.
    # @RELATION: [DEPENDS_ON] ->[ValidationFinding]
    def _upsert_clarification_finding(
        self,
        session: DatasetReviewSession,
        question: ClarificationQuestion,
        answer_kind: AnswerKind,
        answer_value: Optional[str],
    ) -> ValidationFinding:
        caused_by_ref = f"clarification:{question.question_id}"
        existing = next(
            (
                finding
                for finding in session.findings
                if finding.area == FindingArea.CLARIFICATION
                and finding.caused_by_ref == caused_by_ref
            ),
            None,
        )

        if answer_kind in {AnswerKind.SELECTED, AnswerKind.CUSTOM}:
            resolution_state = ResolutionState.RESOLVED
            resolved_at = datetime.utcnow()
            message = f"Clarified '{question.topic_ref}' with answer '{answer_value}'."
        elif answer_kind == AnswerKind.SKIPPED:
            resolution_state = ResolutionState.SKIPPED
            resolved_at = None
            message = f"Clarification for '{question.topic_ref}' was skipped and still needs review."
        else:
            resolution_state = ResolutionState.EXPERT_REVIEW
            resolved_at = None
            message = (
                f"Clarification for '{question.topic_ref}' requires expert review."
            )

        if existing is None:
            existing = ValidationFinding(
                finding_id=str(uuid.uuid4()),
                session_id=session.session_id,
                area=FindingArea.CLARIFICATION,
                severity=FindingSeverity.WARNING,
                code="CLARIFICATION_PENDING",
                title="Clarification pending",
                message=message,
                resolution_state=resolution_state,
                resolution_note=None,
                caused_by_ref=caused_by_ref,
                created_at=datetime.utcnow(),
                resolved_at=resolved_at,
            )
            self.repository.db.add(existing)
            session.findings.append(existing)
        else:
            existing.message = message
            existing.resolution_state = resolution_state
            existing.resolved_at = resolved_at

        if answer_kind in {AnswerKind.SELECTED, AnswerKind.CUSTOM}:
            existing.code = "CLARIFICATION_RESOLVED"
            existing.title = "Clarification resolved"
        elif answer_kind == AnswerKind.SKIPPED:
            existing.code = "CLARIFICATION_SKIPPED"
            existing.title = "Clarification skipped"
        else:
            existing.code = "CLARIFICATION_EXPERT_REVIEW"
            existing.title = "Clarification requires expert review"

        return existing

    # [/DEF:_upsert_clarification_finding:Function]

    # [DEF:_derive_readiness_state:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Recompute readiness after clarification mutation while preserving unresolved visibility semantics.
    # @RELATION: [DEPENDS_ON] ->[ClarificationSession]
    # @RELATION: [DEPENDS_ON] ->[DatasetReviewSession]
    def _derive_readiness_state(self, session: DatasetReviewSession) -> ReadinessState:
        clarification_session = self._get_latest_clarification_session(session)
        if clarification_session is None:
            return session.readiness_state

        if clarification_session.current_question_id:
            return ReadinessState.CLARIFICATION_ACTIVE

        if clarification_session.remaining_count > 0:
            return ReadinessState.CLARIFICATION_NEEDED

        return ReadinessState.REVIEW_READY

    # [/DEF:_derive_readiness_state:Function]

    # [DEF:_derive_recommended_action:Function]
    # @COMPLEXITY: 2
    # @PURPOSE: Recompute next-action guidance after clarification mutations.
    def _derive_recommended_action(
        self, session: DatasetReviewSession
    ) -> RecommendedAction:
        clarification_session = self._get_latest_clarification_session(session)
        if clarification_session is None:
            return session.recommended_action
        if clarification_session.current_question_id:
            return RecommendedAction.ANSWER_NEXT_QUESTION
        if clarification_session.remaining_count > 0:
            return RecommendedAction.START_CLARIFICATION
        return RecommendedAction.REVIEW_DOCUMENTATION

    # [/DEF:_derive_recommended_action:Function]


# [/DEF:ClarificationEngine:Class]

# [/DEF:ClarificationEngine:Module]