"""
State machine implementation for chocolate tempering process control.
Manages the complete tempering process workflow with proper state transitions.
"""

import asyncio
import logging
from typing import Optional, Dict, Any, Callable
from datetime import datetime, timedelta
from enum import Enum
from dataclasses import dataclass

from statemachine import StateMachine, State
from statemachine.exceptions import TransitionNotAllowed

from ...shared.models.recipe import RecipePhase
from ...shared.config import settings
from ..hardware.hardware_manager import hardware_manager


logger = logging.getLogger(__name__)


@dataclass
class ProcessMetrics:
    """Process execution metrics and measurements."""
    phase_start_time: datetime
    target_temperature: Optional[float] = None
    current_temperature: Optional[float] = None
    temperature_error: float = 0.0
    stability_duration: float = 0.0  # Time at stable temperature
    energy_consumption: float = 0.0
    error_count: int = 0
    warning_count: int = 0


@dataclass
class ProcessContext:
    """Context information for process execution."""
    recipe_id: int
    session_id: str
    recipe_parameters: Dict[str, Any]
    current_metrics: ProcessMetrics
    user_id: Optional[str] = None
    start_time: Optional[datetime] = None
    estimated_completion: Optional[datetime] = None


class TemperingStateMachine(StateMachine):
    """
    State machine for chocolate tempering process control.
    Implements the complete tempering workflow with safety checks.
    """
    
    # Define states based on recipe phases
    idle = State("Idle", initial=True)
    preheating = State("Preheating")
    heating = State("Heating")
    heating_delay = State("HeatingDelay")
    cooling = State("Cooling") 
    cooling_delay = State("CoolingDelay")
    pouring = State("Pouring")
    completed = State("Completed")
    error = State("Error")
    emergency_stopped = State("EmergencyStopped")
    
    # Define transitions
    start_process = idle.to(preheating)
    begin_heating = preheating.to(heating)
    start_heating_delay = heating.to(heating_delay)
    begin_cooling = heating_delay.to(cooling)
    start_cooling_delay = cooling.to(cooling_delay)
    begin_pouring = cooling_delay.to(pouring)
    finish_process = pouring.to(completed)
    
    # Error and emergency transitions (from any state)
    error_occurred = (
        preheating.to(error) |
        heating.to(error) |
        heating_delay.to(error) |
        cooling.to(error) |
        cooling_delay.to(error) |
        pouring.to(error)
    )
    
    emergency_stop = (
        preheating.to(emergency_stopped) |
        heating.to(emergency_stopped) |
        heating_delay.to(emergency_stopped) |
        cooling.to(emergency_stopped) |
        cooling_delay.to(emergency_stopped) |
        pouring.to(emergency_stopped) |
        error.to(emergency_stopped)
    )
    
    # Recovery transitions
    reset_from_error = error.to(idle)
    reset_from_emergency = emergency_stopped.to(idle)
    reset_from_completed = completed.to(idle)
    
    def __init__(self, context: ProcessContext):
        self.context = context
        self.phase_timers: Dict[str, datetime] = {}
        self.stability_checker = TemperatureStabilityChecker()
        self.safety_monitor = ProcessSafetyMonitor()
        self.control_loop_task: Optional[asyncio.Task] = None
        self._running = False
        
        super().__init__()
        
        # Set up phase-specific handlers
        self._phase_handlers = {
            "preheating": self._handle_preheating_phase,
            "heating": self._handle_heating_phase,
            "heating_delay": self._handle_heating_delay_phase,
            "cooling": self._handle_cooling_phase,
            "cooling_delay": self._handle_cooling_delay_phase,
            "pouring": self._handle_pouring_phase,
        }
    
    async def initialize(self) -> bool:
        """Initialize the state machine and hardware systems."""
        try:
            logger.info(f"Initializing tempering process for recipe {self.context.recipe_id}")
            
            # Verify hardware is ready
            is_safe, issues = await hardware_manager.is_safe_to_operate()
            if not is_safe:
                logger.error(f"Hardware not safe to operate: {issues}")
                return False
            
            # Initialize process metrics
            self.context.current_metrics = ProcessMetrics(
                phase_start_time=datetime.now(),
                target_temperature=self.context.recipe_parameters.get("heating_goal", 46.0)
            )
            
            # Start safety monitoring
            self.safety_monitor.initialize(self.context)
            
            logger.info("Tempering state machine initialized successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize state machine: {e}")
            return False
    
    async def start_control_loop(self) -> None:
        """Start the main process control loop."""
        if self.control_loop_task and not self.control_loop_task.done():
            return
        
        self._running = True
        self.control_loop_task = asyncio.create_task(self._control_loop())
        logger.info("Process control loop started")
    
    async def stop_control_loop(self) -> None:
        """Stop the main process control loop."""
        self._running = False
        if self.control_loop_task:
            self.control_loop_task.cancel()
            try:
                await self.control_loop_task
            except asyncio.CancelledError:
                pass
        logger.info("Process control loop stopped")
    
    async def _control_loop(self) -> None:
        """Main control loop for process execution."""
        while self._running:
            try:
                current_state = self.current_state.name.lower()
                
                # Execute phase-specific handler
                if current_state in self._phase_handlers:
                    await self._phase_handlers[current_state]()
                
                # Check for safety issues
                await self._check_safety_conditions()
                
                # Update metrics
                await self._update_process_metrics()
                
                # Control loop delay
                await asyncio.sleep(settings.process.process_control_interval)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in control loop: {e}")
                try:
                    self.error_occurred()
                except TransitionNotAllowed:
                    pass
                await asyncio.sleep(1.0)  # Error recovery delay
    
    async def _handle_preheating_phase(self) -> None:
        """Handle preheating phase - warm up hardware."""
        # Get current temperatures
        tank_temp = await hardware_manager.get_average_tank_temperature()
        pump_temp_reading = await hardware_manager.get_temperature("pump")
        pump_temp = pump_temp_reading.value if pump_temp_reading else 0.0
        
        # Check if preheating targets are met
        tank_target = settings.temperature.tank_max_heat
        pump_target = settings.temperature.pump_max_heat
        
        if tank_temp and tank_temp >= tank_target and pump_temp >= pump_target:
            logger.info("Preheating complete, transitioning to heating phase")
            try:
                self.begin_heating()
            except TransitionNotAllowed as e:
                logger.error(f"Cannot transition to heating: {e}")
        else:
            # Enable heaters for preheating
            await hardware_manager.enable_heater("tank_heater")
            await hardware_manager.enable_heater("pump_heater")
            
            # Update metrics
            self.context.current_metrics.current_temperature = tank_temp
            self.context.current_metrics.target_temperature = tank_target
    
    async def _handle_heating_phase(self) -> None:
        """Handle heating phase - reach heating goal temperature."""
        heating_goal = self.context.recipe_parameters.get("heating_goal", 46.0)
        fountain_temp_reading = await hardware_manager.get_temperature("fountain")
        fountain_temp = fountain_temp_reading.value if fountain_temp_reading else 0.0
        
        # Enable motors as specified in recipe
        if self.context.recipe_parameters.get("mixer_enabled", True):
            await hardware_manager.enable_motor("mixer_motor")
        
        if self.context.recipe_parameters.get("fountain_enabled", True):
            await hardware_manager.enable_motor("fountain_motor")
        
        # Temperature control
        temp_error = heating_goal - fountain_temp
        tolerance = settings.temperature.temperature_tolerance
        
        if abs(temp_error) <= tolerance:
            # Temperature reached, check stability
            if self.stability_checker.is_stable(fountain_temp, heating_goal, tolerance):
                logger.info(f"Heating goal reached and stable: {fountain_temp}°C")
                try:
                    self.start_heating_delay()
                except TransitionNotAllowed as e:
                    logger.error(f"Cannot transition to heating delay: {e}")
        else:
            # Adjust heating based on temperature error
            if temp_error > 0:
                # Need more heat
                await hardware_manager.enable_heater("tank_heater")
                await hardware_manager.enable_heater("pump_heater")
            else:
                # Too hot, reduce heating
                await hardware_manager.disable_heater("tank_heater")
        
        # Update metrics
        self.context.current_metrics.current_temperature = fountain_temp
        self.context.current_metrics.target_temperature = heating_goal
        self.context.current_metrics.temperature_error = temp_error
    
    async def _handle_heating_delay_phase(self) -> None:
        """Handle heating delay phase - maintain temperature for specified time."""
        if not self._check_phase_timer("heating_delay", settings.process.heating_delay):
            # Still in delay period, maintain temperature
            await self._maintain_temperature(
                self.context.recipe_parameters.get("heating_goal", 46.0)
            )
        else:
            logger.info("Heating delay complete, transitioning to cooling")
            try:
                self.begin_cooling()
            except TransitionNotAllowed as e:
                logger.error(f"Cannot transition to cooling: {e}")
    
    async def _handle_cooling_phase(self) -> None:
        """Handle cooling phase - cool to cooling goal temperature."""
        cooling_goal = self.context.recipe_parameters.get("cooling_goal", 27.0)
        fountain_temp_reading = await hardware_manager.get_temperature("fountain")
        fountain_temp = fountain_temp_reading.value if fountain_temp_reading else 50.0
        
        # Disable tank heaters for cooling
        await hardware_manager.disable_heater("tank_heater")
        
        # Keep pump heater on to maintain pump temperature
        await hardware_manager.enable_heater("pump_heater")
        
        # Keep motors running for circulation
        if self.context.recipe_parameters.get("mixer_enabled", True):
            await hardware_manager.enable_motor("mixer_motor")
        
        if self.context.recipe_parameters.get("fountain_enabled", True):
            await hardware_manager.enable_motor("fountain_motor")
        
        # Check if cooling goal is reached
        tolerance = settings.temperature.temperature_tolerance
        temp_error = fountain_temp - cooling_goal
        
        if temp_error <= tolerance and fountain_temp <= cooling_goal:
            if self.stability_checker.is_stable(fountain_temp, cooling_goal, tolerance):
                logger.info(f"Cooling goal reached and stable: {fountain_temp}°C")
                try:
                    self.start_cooling_delay()
                except TransitionNotAllowed as e:
                    logger.error(f"Cannot transition to cooling delay: {e}")
        
        # Update metrics
        self.context.current_metrics.current_temperature = fountain_temp
        self.context.current_metrics.target_temperature = cooling_goal
        self.context.current_metrics.temperature_error = temp_error
    
    async def _handle_cooling_delay_phase(self) -> None:
        """Handle cooling delay phase - maintain cooling temperature."""
        if not self._check_phase_timer("cooling_delay", settings.process.cooling_delay):
            # Still in delay period, maintain temperature
            await self._maintain_temperature(
                self.context.recipe_parameters.get("cooling_goal", 27.0)
            )
        else:
            logger.info("Cooling delay complete, transitioning to pouring")
            try:
                self.begin_pouring()
            except TransitionNotAllowed as e:
                logger.error(f"Cannot transition to pouring: {e}")
    
    async def _handle_pouring_phase(self) -> None:
        """Handle pouring phase - maintain pouring temperature and manage pedal control."""
        pouring_goal = self.context.recipe_parameters.get("pouring_goal")
        if not pouring_goal:
            # Use cooling goal if no specific pouring goal
            pouring_goal = self.context.recipe_parameters.get("cooling_goal", 30.0)
        
        # Maintain pouring temperature
        await self._maintain_temperature(pouring_goal)
        
        # Handle pedal control if enabled
        if self.context.recipe_parameters.get("pedal_control_enabled", True):
            await self._manage_pedal_control()
        
        # Keep circulation motors running
        if self.context.recipe_parameters.get("fountain_enabled", True):
            await hardware_manager.enable_motor("fountain_motor")
        
        # Update metrics
        fountain_temp_reading = await hardware_manager.get_temperature("fountain")
        fountain_temp = fountain_temp_reading.value if fountain_temp_reading else pouring_goal
        
        self.context.current_metrics.current_temperature = fountain_temp
        self.context.current_metrics.target_temperature = pouring_goal
        self.context.current_metrics.temperature_error = fountain_temp - pouring_goal
        
        # Pouring phase continues until manually stopped or error occurs
        # In a real system, this might be triggered by operator action or timer
    
    async def _maintain_temperature(self, target_temp: float) -> None:
        """Maintain target temperature with PID control."""
        fountain_temp_reading = await hardware_manager.get_temperature("fountain")
        current_temp = fountain_temp_reading.value if fountain_temp_reading else target_temp
        
        temp_error = current_temp - target_temp
        tolerance = settings.temperature.temperature_tolerance
        
        if abs(temp_error) > tolerance:
            if temp_error < 0:
                # Need heating
                await hardware_manager.enable_heater("pump_heater")
                if temp_error < -2.0:  # Significant heating needed
                    await hardware_manager.enable_heater("tank_heater")
            else:
                # Too hot, reduce heating
                await hardware_manager.disable_heater("tank_heater")
                if temp_error > 2.0:  # Significant cooling needed
                    await hardware_manager.disable_heater("pump_heater")
    
    async def _manage_pedal_control(self) -> None:
        """Manage automatic pedal control for pouring."""
        # This would implement automatic pedal cycling based on recipe parameters
        # For now, just log the action
        on_time = self.context.recipe_parameters.get("pedal_on_time", 2.0)
        off_time = self.context.recipe_parameters.get("pedal_off_time", 3.0)
        
        logger.debug(f"Pedal control active: {on_time}s on, {off_time}s off")
        # Implementation would control pedal solenoid via hardware manager
    
    def _check_phase_timer(self, phase_name: str, duration_seconds: float) -> bool:
        """Check if a phase timer has completed."""
        if phase_name not in self.phase_timers:
            self.phase_timers[phase_name] = datetime.now()
            return False
        
        elapsed = (datetime.now() - self.phase_timers[phase_name]).total_seconds()
        return elapsed >= duration_seconds
    
    async def _check_safety_conditions(self) -> None:
        """Check safety conditions and trigger emergency stop if needed."""
        is_safe, issues = await hardware_manager.is_safe_to_operate()
        
        if not is_safe:
            logger.error(f"Safety issue detected: {issues}")
            try:
                self.emergency_stop()
            except TransitionNotAllowed:
                pass
    
    async def _update_process_metrics(self) -> None:
        """Update process execution metrics."""
        self.context.current_metrics.phase_start_time = datetime.now()
        
        # Update stability duration if temperature is stable
        if hasattr(self, 'stability_checker'):
            current_temp = self.context.current_metrics.current_temperature
            target_temp = self.context.current_metrics.target_temperature
            
            if (current_temp is not None and target_temp is not None and 
                abs(current_temp - target_temp) <= settings.temperature.temperature_tolerance):
                # Temperature is stable, update duration
                pass
    
    def get_current_phase(self) -> RecipePhase:
        """Get current recipe phase."""
        state_to_phase = {
            "idle": RecipePhase.COMPLETED,
            "preheating": RecipePhase.PREHEATING,
            "heating": RecipePhase.HEATING,
            "heating_delay": RecipePhase.HEATING_DELAY,
            "cooling": RecipePhase.COOLING,
            "cooling_delay": RecipePhase.COOLING_DELAY,
            "pouring": RecipePhase.POURING,
            "completed": RecipePhase.COMPLETED,
            "error": RecipePhase.ERROR,
            "emergency_stopped": RecipePhase.STOPPED,
        }
        return state_to_phase.get(self.current_state.name.lower(), RecipePhase.ERROR)
    
    def get_process_status(self) -> Dict[str, Any]:
        """Get comprehensive process status."""
        return {
            "current_state": self.current_state.name,
            "current_phase": self.get_current_phase().value,
            "recipe_id": self.context.recipe_id,
            "session_id": self.context.session_id,
            "start_time": self.context.start_time.isoformat() if self.context.start_time else None,
            "current_temperature": self.context.current_metrics.current_temperature,
            "target_temperature": self.context.current_metrics.target_temperature,
            "temperature_error": self.context.current_metrics.temperature_error,
            "error_count": self.context.current_metrics.error_count,
            "warning_count": self.context.current_metrics.warning_count,
            "is_running": self._running,
            "phase_timers": {k: v.isoformat() for k, v in self.phase_timers.items()},
        }


class TemperatureStabilityChecker:
    """Utility class to check temperature stability."""
    
    def __init__(self, stability_time: float = 30.0):
        self.stability_time = stability_time  # seconds
        self.temperature_history: Dict[str, list] = {}
    
    def is_stable(self, current_temp: float, target_temp: float, tolerance: float) -> bool:
        """Check if temperature has been stable within tolerance for required time."""
        key = f"{target_temp}"
        now = datetime.now()
        
        # Initialize history if needed
        if key not in self.temperature_history:
            self.temperature_history[key] = []
        
        history = self.temperature_history[key]
        
        # Add current reading
        history.append((now, current_temp))
        
        # Remove old readings
        cutoff_time = now - timedelta(seconds=self.stability_time)
        history[:] = [(time, temp) for time, temp in history if time > cutoff_time]
        
        # Check if all recent readings are within tolerance
        if len(history) < 3:  # Need at least a few readings
            return False
        
        return all(abs(temp - target_temp) <= tolerance for _, temp in history)


class ProcessSafetyMonitor:
    """Monitor process safety during execution."""
    
    def __init__(self):
        self.context: Optional[ProcessContext] = None
        self.last_safety_check = datetime.now()
    
    def initialize(self, context: ProcessContext) -> None:
        """Initialize safety monitor with process context."""
        self.context = context
    
    async def check_safety(self) -> tuple[bool, list[str]]:
        """Perform comprehensive safety check."""
        issues = []
        
        # Check hardware safety
        is_hardware_safe, hardware_issues = await hardware_manager.is_safe_to_operate()
        if not is_hardware_safe:
            issues.extend(hardware_issues)
        
        # Check temperature limits
        for sensor_name, reading in (await hardware_manager.get_all_temperatures()).items():
            if not reading.is_valid:
                issues.append(f"Temperature sensor {sensor_name}: {reading.error_message}")
        
        # Update last check time
        self.last_safety_check = datetime.now()
        
        return len(issues) == 0, issues