# File: src/noveler/infrastructure/services/unified_session_executor.py
# Purpose: Orchestrates unified writing sessions across five and ten stage workflows.
# Context: Bridges legacy Claude pipelines with the newer A30 detailed execution model.
#         Uses content_cleaner, turn_allocation_calculator, and legacy_format_converter for modularity.
"""Unified session executor bridging five-stage and ten-stage writing flows."""

from __future__ import annotations

import asyncio
import uuid
from collections.abc import Mapping, MutableMapping
from dataclasses import dataclass, field
from typing import Any

from noveler.application.services.a30_compatibility_adapter import (
    A30CompatibilityAdapter,
    CompatibilityMode,
)
from noveler.application.use_cases.ten_stage_progress_use_case import (
    TenStageProgressRequest,
    TenStageProgressUseCase,
)
from noveler.domain.value_objects.detailed_execution_stage import DetailedExecutionStage
from noveler.domain.value_objects.five_stage_writing_execution import (
    ExecutionStage,
    FiveStageWritingRequest,
)
from noveler.infrastructure.logging.unified_logger import get_logger
from noveler.infrastructure.services.content_cleaner import ContentCleaner
from noveler.infrastructure.services.content_validation_service import ContentValidationService
from noveler.infrastructure.services.legacy_format_converter import LegacyFormatConverter
from noveler.infrastructure.services.turn_allocation_calculator import (
    TurnAllocation,
    TurnAllocationCalculator,
)

StageKey = ExecutionStage | DetailedExecutionStage


@dataclass
class SessionState:
    """Lightweight state tracker used while a unified writing session executes."""

    session_id: str
    total_turns_available: int = 30
    turns_used: int = 0
    error_count: int = 0
    current_stage: StageKey | None = None
    stage_outputs: dict[StageKey, str] = field(default_factory=dict)
    accumulated_context: str = ""
    metadata: dict[str, Any] = field(default_factory=dict)
    use_ten_stage: bool = False

    @property
    def remaining_turns(self) -> int:
        """Return the number of turns that can still be consumed."""
        remaining = self.total_turns_available - self.turns_used
        return remaining if remaining > 0 else 0

    def can_continue(self) -> bool:
        """Determine whether the executor should keep processing stages."""
        return self.remaining_turns > 0 and self.error_count < 3


@dataclass
class UnifiedSessionResponse:
    """Structured response returned after a unified session completes."""

    session_id: str
    status: str
    stage_outputs: Mapping[StageKey, str] = field(default_factory=dict)
    turns_used: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)
    raw_responses: list[dict[str, Any]] = field(default_factory=list)
    use_ten_stage: bool = False
    legacy_stage_outputs: Mapping[ExecutionStage, str] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Return a serialisable view of the response."""
        stage_payload = {
            _stage_key_to_value(stage): output
            for stage, output in self.stage_outputs.items()
        }
        legacy_payload = {
            stage.value: output for stage, output in self.legacy_stage_outputs.items()
        }
        return {
            "session_id": self.session_id,
            "status": self.status,
            "turns_used": self.turns_used,
            "stage_outputs": stage_payload,
            "metadata": self.metadata,
            "raw_responses": self.raw_responses,
            "use_ten_stage": self.use_ten_stage,
            "legacy_stage_outputs": legacy_payload,
        }


class UnifiedSessionExecutor:
    """Coordinates unified writing sessions across multiple execution modes."""

    def __init__(
        self,
        claude_service: Any = None,
        validation_service: ContentValidationService | None = None,
        *,
        compatibility_adapter: A30CompatibilityAdapter | None = None,
        progress_use_case: TenStageProgressUseCase | None = None,
        logger_name: str = __name__,
        total_turn_limit: int = 30,
    ) -> None:
        """Initialize the unified session executor.

        Args:
            claude_service: The Claude service for executing stages.
            validation_service: Optional content validation service.
            compatibility_adapter: Optional A30 compatibility adapter.
            progress_use_case: Optional use case for tracking progress.
            logger_name: Logger name (defaults to module name).
            total_turn_limit: Maximum turns available per session (defaults to 30).

        Raises:
            ValueError: If claude_service is not provided.
        """
        if claude_service is None:
            raise ValueError("claude_service is required")
        self.claude_service = claude_service
        self.validation_service = validation_service or ContentValidationService()
        self.compatibility_adapter = compatibility_adapter
        self.progress_use_case = progress_use_case
        self.total_turn_limit = total_turn_limit
        self._logger = get_logger(logger_name)

        # Initialize sub-components for modularity
        self._allocator = TurnAllocationCalculator()
        self._content_cleaner = ContentCleaner()
        self._legacy_converter = LegacyFormatConverter(compatibility_adapter)

        # Expose default allocations for backward compatibility with tests
        self.default_allocations_5stage = self._allocator.default_allocations_5stage
        self.default_allocations_10stage = self._allocator.default_allocations_10stage
        self.default_allocations = self.default_allocations_5stage

    async def execute_unified_session(self, request: FiveStageWritingRequest) -> UnifiedSessionResponse:
        """Execute a unified writing session (5-stage or 10-stage).

        Args:
            request: The writing request containing stage, episode, and turn information.

        Returns:
            UnifiedSessionResponse with stage outputs, turn usage, and metadata.
        """
        use_ten_stage = self._determine_execution_mode(request)
        allocations = (
            self._allocator.calculate_turn_allocations_10stage(request)
            if use_ten_stage
            else self._calculate_turn_allocations(request)
        )

        session_id = self._generate_session_id(request)
        state = SessionState(
            session_id=session_id,
            total_turns_available=self.total_turn_limit,
            use_ten_stage=use_ten_stage,
        )

        raw_responses: list[dict[str, Any]] = []
        context = state.accumulated_context

        if use_ten_stage and self.progress_use_case:
            await self._safe_progress_call(
                TenStageProgressRequest(
                    episode_number=request.episode_number,
                    project_root=request.project_root,
                    operation="start",
                )
            )

        for stage, allocation in allocations.items():
            if not state.can_continue():
                break

            state.current_stage = stage
            remaining_turns = state.remaining_turns
            if remaining_turns <= 0:
                break

            max_turns = min(allocation.max_turns, remaining_turns)
            try:
                response = await self.claude_service.execute_with_turn_limit(
                    stage=stage,
                    request=request,
                    max_turns=max_turns,
                    context=context,
                )
            except Exception as exc:  # pragma: no cover - defensive logging path
                state.error_count += 1
                self._logger.exception(
                    "Stage execution failed", extra={"stage": _stage_key_to_value(stage), "error": str(exc)}
                )
                if not state.can_continue():
                    break
                continue

            raw_responses.append({"stage": _stage_key_to_value(stage), "response": response})

            turns_used = int(response.get("turns_used") or 0)
            state.turns_used += turns_used

            output_text = response.get("output") or ""
            cleaned_output = await self._validate_and_clean_output(output_text, stage)
            state.stage_outputs[stage] = cleaned_output

            if cleaned_output:
                context = f"{context}\n{cleaned_output}".strip()
                state.accumulated_context = context

            if not response.get("success", True):
                state.error_count += 1
                if not state.can_continue():
                    break

            if use_ten_stage and self.progress_use_case:
                await self._safe_progress_call(
                    TenStageProgressRequest(
                        episode_number=request.episode_number,
                        project_root=request.project_root,
                        operation="update",
                        stage=stage if isinstance(stage, DetailedExecutionStage) else None,
                    )
                )

            if state.turns_used >= state.total_turns_available:
                break

        status = "completed" if state.error_count < 3 else "failed"
        if state.turns_used >= state.total_turns_available and status == "completed":
            status = "completed_with_turn_limit"

        metadata: dict[str, Any] = {
            "execution_mode": "10-stage" if use_ten_stage else "5-stage",
            "turn_limit": state.total_turns_available,
            "turns_used": state.turns_used,
            "stages_completed": len(state.stage_outputs),
        }
        metadata["unified_session"] = True
        if use_ten_stage:
            metadata["total_detailed_stages"] = len(DetailedExecutionStage)

        legacy_outputs: Mapping[ExecutionStage, str] = {}
        if use_ten_stage and state.stage_outputs:
            legacy_outputs = self._convert_to_legacy_format(
                {stage: output for stage, output in state.stage_outputs.items() if isinstance(stage, DetailedExecutionStage)}
            )

        return UnifiedSessionResponse(
            session_id=state.session_id,
            status=status,
            stage_outputs=dict(state.stage_outputs),
            turns_used=state.turns_used,
            metadata=metadata,
            raw_responses=raw_responses,
            use_ten_stage=use_ten_stage,
            legacy_stage_outputs=legacy_outputs,
        )

    def get_stage_allocations(self, request: FiveStageWritingRequest) -> Mapping[StageKey, TurnAllocation]:
        """Return an ordered mapping of stage allocations for the request.

        Args:
            request: The writing request to get allocations for.

        Returns:
            Mapping of stages to turn allocations.
        """
        if self._determine_execution_mode(request):
            return self._allocator.calculate_turn_allocations_10stage(request)
        return self._allocator.calculate_turn_allocations_5stage(request)

    def _determine_execution_mode(self, request: FiveStageWritingRequest) -> bool:
        """Determine if the current session should run in ten-stage mode.

        Checks in order:
        1. Explicit request.use_ten_stage attribute
        2. Compatibility adapter mode (if available)
        3. Default to ten-stage mode

        Args:
            request: The writing request to check.

        Returns:
            True for ten-stage mode, False for five-stage mode.
        """
        explicit_choice = getattr(request, "use_ten_stage", None)
        if explicit_choice is not None:
            return bool(explicit_choice)

        if self.compatibility_adapter is None:
            return False

        mode = self.compatibility_adapter.compatibility_mode
        if mode == CompatibilityMode.A30_DETAILED_TEN_STAGE:
            return True
        if mode == CompatibilityMode.LEGACY_FIVE_STAGE:
            return False

        return True

    async def _validate_and_clean_output(self, content: str, stage: StageKey) -> str:
        """Validate and sanitise stage output before storing it.

        Args:
            content: The raw content from Claude service.
            stage: The execution stage for validation context.

        Returns:
            Cleaned and validated content.
        """
        cleaned = self._content_cleaner.clean_all(content)

        if not cleaned.strip():
            cleaned = content.strip()

        validation_result = None
        if self.validation_service:
            try:
                result_or_coro = self.validation_service.validate(cleaned, _stage_key_to_value(stage))
                validation_result = await _ensure_awaited(result_or_coro)
            except Exception:  # pragma: no cover - validation is best effort
                self._logger.debug("Validation service raised an exception", exc_info=True)

        if validation_result and getattr(validation_result, "cleaned_content", None):
            cleaned = validation_result.cleaned_content

        return cleaned.strip()

    def _convert_to_legacy_format(
        self, detailed_outputs: Mapping[DetailedExecutionStage, str]
    ) -> Mapping[ExecutionStage, str]:
        """Combine detailed-stage outputs into their legacy equivalents.

        Args:
            detailed_outputs: Mapping of detailed stages to their outputs.

        Returns:
            Mapping of legacy stages to combined outputs.
        """
        return self._legacy_converter.convert_to_legacy_format(detailed_outputs)

    async def _safe_progress_call(self, request: TenStageProgressRequest) -> None:
        """Call the progress use case and suppress non-critical failures.

        Args:
            request: The progress tracking request.
        """
        if not self.progress_use_case:
            return
        try:
            result = self.progress_use_case.execute(request)
            await _ensure_awaited(result)
        except Exception:  # pragma: no cover - telemetry only
            self._logger.debug("Progress tracking call failed", exc_info=True)

    def _generate_session_id(self, request: FiveStageWritingRequest) -> str:
        """Generate a stable session identifier for the request.

        Args:
            request: The writing request containing episode information.

        Returns:
            Session identifier compatible with legacy expectations.
        """
        existing = getattr(request, "session_id", None)
        if existing:
            return str(existing)

        episode = getattr(request, "episode_number", 0) or 0
        return f"unified_{int(episode):03d}_{uuid.uuid4().hex[:8]}"

    def _estimate_complexity(self, request: FiveStageWritingRequest) -> float:
        """Delegate complexity estimation to the allocation calculator.

        Args:
            request: The writing request to evaluate.

        Returns:
            Complexity factor between 0.0 and 1.0.
        """
        return self._allocator._estimate_complexity(request)

    def _calculate_turn_allocations(self, request: FiveStageWritingRequest) -> MutableMapping[StageKey, TurnAllocation]:
        """Return the turn allocations for the request respecting execution mode.

        Args:
            request: The writing request describing episode parameters.

        Returns:
            Mapping of stages to their allocations.
        """
        if self._determine_execution_mode(request):
            return self._allocator.calculate_turn_allocations_10stage(request)

        allocations = dict(self._allocator.calculate_turn_allocations_5stage(request))
        total_max = sum(allocation.max_turns for allocation in allocations.values())
        overflow = max(0, total_max - self.total_turn_limit)

        if overflow <= 0:
            return allocations

        adjustable = sorted(
            allocations.values(),
            key=lambda entry: (entry.priority, entry.max_turns),
            reverse=True,
        )
        for allocation in adjustable:
            if overflow <= 0:
                break
            reducible = max(0, allocation.max_turns - allocation.min_turns)
            if reducible <= 0:
                continue
            reduction = min(reducible, overflow)
            allocation.max_turns -= reduction
            overflow -= reduction

        return allocations

    def _remove_json_metadata(self, content: str) -> str:
        """Remove JSON artefacts from content using the cleaner pipeline.

        Args:
            content: The content to process.

        Returns:
            Content without JSON metadata blocks.
        """
        return self._content_cleaner.remove_json_metadata(content)

    def _remove_prompt_contamination(self, content: str) -> str:
        """Strip prompt scaffolding that leaked into the output.

        Args:
            content: Model output potentially containing prompt text.

        Returns:
            Content with prompt contamination removed.
        """
        return self._content_cleaner.remove_prompt_contamination(content)

    def _extract_manuscript_text(self, content: str) -> str:
        """Extract manuscript body starting from the first numbered heading.

        Args:
            content: The content to extract manuscript text from.

        Returns:
            Manuscript text beginning with the numbered heading when present.
        """
        return self._content_cleaner.extract_manuscript_text(content)

    def _auto_clean_content(self, content: str) -> str:
        """Apply the full cleaning pipeline and return the manuscript body.

        Args:
            content: Raw stage output text.

        Returns:
            Cleaned manuscript content ready for validation.
        """
        cleaned = self._content_cleaner.clean_all(content)
        return self._content_cleaner.extract_manuscript_text(cleaned)


def _stage_key_to_value(stage: StageKey) -> str:
    """Return a stable string representation for either stage type.

    Args:
        stage: The execution stage (5-stage or 10-stage).

    Returns:
        String representation of the stage.
    """
    if hasattr(stage, "value"):
        return stage.value  # type: ignore[return-value]
    if hasattr(stage, "name"):
        return stage.name  # type: ignore[return-value]
    return str(stage)


async def _ensure_awaited(value: Any) -> Any:
    """Await the value when it behaves like a coroutine.

    Args:
        value: The value to potentially await.

    Returns:
        The awaited value or the original value.
    """
    if asyncio.iscoroutine(value) or hasattr(value, "__await__"):
        return await value
    return value
