"""
Data models for iterative reasoning.

This module contains the core data structures used throughout the Think service's
iterative reasoning process.
"""

from dataclasses import dataclass, asdict
from datetime import datetime
from enum import Enum
from typing import Dict, Any, Optional, List


class StepAction(Enum):
    """Available reasoning step actions"""
    CALL_SERVICE = "call_service"
    ANALYZE_DATA = "analyze_data"
    SYNTHESIZE_FINAL = "synthesize_final_response"
    CHECK_GOAL_SATISFACTION = "check_goal_satisfaction"


@dataclass
class ReasoningStep:
    """Represents a single step in the iterative reasoning process"""
    action: str  # StepAction value
    target: Optional[str] = None  # Service name or data target
    reasoning: str = ""  # Why this step is needed
    ready: bool = False  # Terminal signal for synthesis

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class StepResult:
    """Result of executing a reasoning step"""
    step: ReasoningStep
    success: bool
    result_data: Any = None
    error_message: str = ""
    execution_time_ms: float = 0
    timestamp: str = ""

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.utcnow().isoformat()

    def to_dict(self) -> Dict[str, Any]:
        result = asdict(self)
        result['step'] = self.step.to_dict()
        return result


class IterativeContext:
    """Rich context for iterative reasoning process"""

    def __init__(self, original_message: str, identity: str, channel: str, modality: str):
        self.original_message = original_message
        self.identity = identity
        self.channel = channel
        self.modality = modality
        self.start_time = datetime.utcnow()

        # Tracking
        self.completed_steps: List[StepResult] = []
        self.failed_steps: List[StepResult] = []
        self.accumulated_knowledge: Dict[str, Any] = {}
        self.service_call_counts: Dict[str, int] = {}
        self.step_count = 0

        # State flags
        self.goal_satisfied = False
        self.force_synthesis = False
        self.max_steps_reached = False
        self.timeout_reached = False

    def add_step_result(self, step_result: StepResult):
        """Add a completed step result"""
        self.step_count += 1

        if step_result.success:
            self.completed_steps.append(step_result)

            # Track service calls
            if step_result.step.action == StepAction.CALL_SERVICE.value:
                service = step_result.step.target
                self.service_call_counts[service] = self.service_call_counts.get(service, 0) + 1

            # Accumulate knowledge from successful steps
            if step_result.result_data:
                step_key = f"step_{self.step_count}_{step_result.step.action}"
                self.accumulated_knowledge[step_key] = step_result.result_data
        else:
            self.failed_steps.append(step_result)

    def get_recent_steps(self, window: int = 3) -> List[StepResult]:
        """Get the last N completed steps"""
        return self.completed_steps[-window:] if len(self.completed_steps) >= window else self.completed_steps

    def has_new_information_in_recent_steps(self, window: int = 3) -> bool:
        """Check if recent steps added meaningful new information"""
        recent_steps = self.get_recent_steps(window)

        if len(recent_steps) < window:
            return True  # Not enough steps to judge

        # Simple heuristic: check if recent steps have substantially different result data
        recent_data_sizes = [len(str(step.result_data)) if step.result_data else 0 for step in recent_steps]

        # If all recent steps produced very little data, might be converging
        if all(size < 100 for size in recent_data_sizes):
            return False

        # Check for diversity in step types
        recent_actions = [step.step.action for step in recent_steps]
        if len(set(recent_actions)) == 1:  # All same action type
            return False

        return True

    def should_stop(self, max_steps: int = 10, max_service_calls: int = 3, max_time_minutes: int = 2) -> tuple[bool, str]:
        """Check all stopping criteria"""

        # Hard limits
        if self.step_count >= max_steps:
            self.max_steps_reached = True
            return True, f"Maximum steps reached ({max_steps})"

        elapsed_minutes = (datetime.utcnow() - self.start_time).total_seconds() / 60
        if elapsed_minutes >= max_time_minutes:
            self.timeout_reached = True
            return True, f"Maximum time reached ({max_time_minutes} minutes)"

        # Service call limits
        for service, count in self.service_call_counts.items():
            if count >= max_service_calls:
                return True, f"Too many calls to {service} service ({count})"

        # Goal satisfaction
        if self.goal_satisfied:
            return True, "Goal satisfaction confirmed"

        # Force synthesis
        if self.force_synthesis:
            return True, "Forced synthesis due to convergence"

        # Convergence detection
        if self.step_count >= 6 and not self.has_new_information_in_recent_steps(3):
            self.force_synthesis = True
            return True, "No new information in recent steps - converged"

        return False, ""

    def to_oracle_context(self) -> Dict[str, Any]:
        """Prepare context for Oracle decision making"""
        return {
            "original_message": self.original_message,
            "step_count": self.step_count,
            "completed_steps": [step.to_dict() for step in self.completed_steps[-5:]],  # Last 5 steps
            "accumulated_knowledge": self.accumulated_knowledge,
            "service_call_counts": self.service_call_counts,
            "failed_steps": [step.to_dict() for step in self.failed_steps[-3:]],  # Last 3 failures
            "goal_satisfied": self.goal_satisfied,
        }

    def get_summary(self) -> str:
        """Get a human-readable summary of the reasoning process"""
        elapsed_time = (datetime.utcnow() - self.start_time).total_seconds()

        summary = f"**Reasoning Summary**\n"
        summary += f"• Steps completed: {len(self.completed_steps)}\n"
        summary += f"• Failed steps: {len(self.failed_steps)}\n"
        summary += f"• Service calls: {dict(self.service_call_counts)}\n"
        summary += f"• Elapsed time: {elapsed_time:.1f}s\n"

        if self.completed_steps:
            summary += f"• Recent actions: {', '.join([s.step.action for s in self.completed_steps[-3:]])}\n"

        return summary