# File: src/noveler/infrastructure/services/turn_allocation_calculator.py
# Purpose: Calculate turn allocations for five-stage and ten-stage writing workflows.
# Context: Extracted from unified_session_executor to separate allocation logic from orchestration.
"""Turn allocation calculator for writing workflows."""

from __future__ import annotations

from collections.abc import MutableMapping
from dataclasses import dataclass

from noveler.domain.value_objects.detailed_execution_stage import DetailedExecutionStage
from noveler.domain.value_objects.five_stage_writing_execution import (
    ExecutionStage,
    FiveStageWritingRequest,
)

StageKey = ExecutionStage | DetailedExecutionStage


@dataclass
class TurnAllocation:
    """Allocation rule describing how many turns a stage may consume."""

    stage: StageKey
    min_turns: int
    max_turns: int
    priority: int
    actual_turns: int = 0


class TurnAllocationCalculator:
    """Calculates turn allocations for writing stages based on complexity and defaults."""

    def __init__(
        self,
        default_allocations_5stage: MutableMapping[ExecutionStage, TurnAllocation] | None = None,
        default_allocations_10stage: MutableMapping[DetailedExecutionStage, TurnAllocation] | None = None,
    ) -> None:
        """Initialize the calculator with default allocations.

        Args:
            default_allocations_5stage: Default allocations for 5-stage workflow.
            default_allocations_10stage: Default allocations for 10-stage workflow.
        """
        self.default_allocations_5stage = default_allocations_5stage or self._build_default_allocations_5stage()
        self.default_allocations_10stage = default_allocations_10stage or self._build_default_allocations_10stage()

    @staticmethod
    def _build_default_allocations_5stage() -> MutableMapping[ExecutionStage, TurnAllocation]:
        """Return default allocations for the legacy five-stage workflow.

        Returns:
            Mapping of execution stages to their default allocations.
        """
        defaults: dict[ExecutionStage, TurnAllocation] = {
            ExecutionStage.DATA_COLLECTION: TurnAllocation(
                ExecutionStage.DATA_COLLECTION, min_turns=3, max_turns=5, priority=3
            ),
            ExecutionStage.PLOT_ANALYSIS: TurnAllocation(
                ExecutionStage.PLOT_ANALYSIS, min_turns=3, max_turns=6, priority=4
            ),
            ExecutionStage.EPISODE_DESIGN: TurnAllocation(
                ExecutionStage.EPISODE_DESIGN, min_turns=4, max_turns=6, priority=4
            ),
            ExecutionStage.MANUSCRIPT_WRITING: TurnAllocation(
                ExecutionStage.MANUSCRIPT_WRITING, min_turns=6, max_turns=8, priority=5
            ),
            ExecutionStage.QUALITY_FINALIZATION: TurnAllocation(
                ExecutionStage.QUALITY_FINALIZATION, min_turns=3, max_turns=5, priority=4
            ),
        }
        return defaults

    @staticmethod
    def _build_default_allocations_10stage() -> MutableMapping[DetailedExecutionStage, TurnAllocation]:
        """Return default allocations for the detailed ten-stage workflow.

        Returns:
            Mapping of detailed stages to their default allocations.
        """
        priority_map: dict[DetailedExecutionStage, int] = {
            DetailedExecutionStage.DATA_COLLECTION: 3,
            DetailedExecutionStage.PLOT_ANALYSIS: 4,
            DetailedExecutionStage.LOGIC_VERIFICATION: 4,
            DetailedExecutionStage.CHARACTER_CONSISTENCY: 4,
            DetailedExecutionStage.DIALOGUE_DESIGN: 3,
            DetailedExecutionStage.EMOTION_CURVE: 3,
            DetailedExecutionStage.SCENE_ATMOSPHERE: 3,
            DetailedExecutionStage.MANUSCRIPT_WRITING: 5,
            DetailedExecutionStage.QUALITY_FINALIZATION: 4,
        }
        allocations: dict[DetailedExecutionStage, TurnAllocation] = {}
        for stage in DetailedExecutionStage:
            min_turns = max(1, stage.expected_turns - 1)
            max_turns = stage.expected_turns + 1
            allocations[stage] = TurnAllocation(stage, min_turns, max_turns, priority_map.get(stage, 3))
        return allocations

    def calculate_turn_allocations_5stage(
        self, request: FiveStageWritingRequest
    ) -> MutableMapping[ExecutionStage, TurnAllocation]:
        """Calculate turn allocations for 5-stage workflow.

        Args:
            request: The writing request with complexity parameters.

        Returns:
            Mapping of stages to adjusted allocations.
        """
        complexity = self._estimate_complexity(request)
        allocations: dict[ExecutionStage, TurnAllocation] = {}
        for stage, default in self.default_allocations_5stage.items():
            extra = max(0, round((default.max_turns - default.min_turns) * complexity))
            max_turns = min(default.max_turns + extra, default.min_turns + 4)
            allocations[stage] = TurnAllocation(stage, default.min_turns, max_turns, default.priority)
        return allocations

    def calculate_turn_allocations_10stage(
        self, request: FiveStageWritingRequest
    ) -> MutableMapping[DetailedExecutionStage, TurnAllocation]:
        """Calculate turn allocations for 10-stage workflow.

        Args:
            request: The writing request with complexity parameters.

        Returns:
            Mapping of detailed stages to adjusted allocations.
        """
        complexity = self._estimate_complexity(request)
        allocations: dict[DetailedExecutionStage, TurnAllocation] = {}
        for stage, default in self.default_allocations_10stage.items():
            flex_window = max(1, round(stage.expected_turns * 0.5))
            bonus = max(0, round(flex_window * complexity))
            max_turns = min(default.max_turns + bonus, stage.max_turns)
            allocations[stage] = TurnAllocation(stage, default.min_turns, max_turns, default.priority)
        return allocations

    @staticmethod
    def _estimate_complexity(request: FiveStageWritingRequest) -> float:
        """Estimate request complexity to adjust turn allocations.

        Args:
            request: The writing request to estimate complexity for.

        Returns:
            Complexity factor between 0.0 and 1.0.
        """
        base = 0.3
        episode_factor = min(request.episode_number / 100.0, 0.4)
        base += episode_factor

        word_target = getattr(request, "word_count_target", 0) or 0
        if word_target > 4000:
            base += 0.1
        if getattr(request, "debug_mode", False):
            base += 0.1
        if getattr(request, "dry_run", False):
            base = max(0.0, base - 0.1)

        return max(0.0, min(1.0, base))
