import os
import pickle
import logging
import time
from datetime import datetime
from typing import Dict, Any, List, Set

logger = logging.getLogger(__name__)

STATE_FILE = "rider_active_state.pkl"

class RiderStateManager:
    """
    Manages the 'Short-Term' Active State of Riders for session persistence.
    Tracks:
    - Minutes Committed (Remaining Workload)
    - Active Kitchens (Unique Pickups in current queue)
    - Last Planned Drop Location (for Daisy Chaining)
    - Timestamp of last update (for Time Decay)
    """
    def __init__(self, state_file: str = STATE_FILE):
        self.state_file = state_file
        self.states = self._load_states()

    def _load_states(self) -> Dict[str, Any]:
        """Load states from pickle."""
        if not os.path.exists(self.state_file):
            return {}
        try:
            with open(self.state_file, 'rb') as f:
                return pickle.load(f)
        except Exception as e:
            logger.error(f"Failed to load rider active states: {e}")
            return {}

    def _save_states(self):
        """Save states to pickle."""
        try:
            with open(self.state_file, 'wb') as f:
                pickle.dump(self.states, f)
        except Exception as e:
            logger.error(f"Failed to save rider active states: {e}")

    def get_rider_state(self, rider_id: int) -> Dict[str, Any]:
        """
        Get the current active state of a rider with TIME DECAY applied.
        If the server restarts after 30 mins, the 'minutes_committed' should reduce by 30.
        """
        rider_id = int(rider_id)
        raw_state = self.states.get(rider_id)
        
        if not raw_state:
            return {
                'minutes_remaining': 0.0,
                'last_drop_lat': None,
                'last_drop_lon': None,
                'active_kitchens': set(),
                'last_updated_ts': time.time()
            }
        
        # Apply Time Decay
        last_ts = raw_state.get('last_updated_ts', time.time())
        current_ts = time.time()
        elapsed_mins = (current_ts - last_ts) / 60.0
        
        remaining = max(0.0, raw_state.get('minutes_remaining', 0.0) - elapsed_mins)
        
        # If queue is empty, kitchens are cleared
        kitchens = raw_state.get('active_kitchens', set())
        if remaining <= 5.0: # Buffer: if almost done, free up kitchens
            kitchens = set()
            
        return {
            'minutes_remaining': remaining,
            'last_drop_lat': raw_state.get('last_drop_lat'),
            'last_drop_lon': raw_state.get('last_drop_lon'),
            'active_kitchens': kitchens,
            'last_updated_ts': current_ts
        }

    def update_rider_state(self, rider_id: int, added_minutes: float, new_kitchens: Set[str], last_lat: float, last_lon: float):
        """
        Update the state after a new assignment.
        """
        rider_id = int(rider_id)
        
        # Get current state (decayed)
        current = self.get_rider_state(rider_id)
        
        # Accumulate
        updated_minutes = current['minutes_remaining'] + added_minutes
        updated_kitchens = current['active_kitchens'].union(new_kitchens)
        
        self.states[rider_id] = {
            'minutes_remaining': updated_minutes,
            'last_drop_lat': last_lat,
            'last_drop_lon': last_lon,
            'active_kitchens': updated_kitchens,
            'last_updated_ts': time.time()
        }
        
        self._save_states()

    def clear_state(self, rider_id: int):
        rider_id = int(rider_id)
        if rider_id in self.states:
            del self.states[rider_id]
            self._save_states()