"""
Cross-month step change state management for persistent anomalies.

Manages step up/down state across month boundaries to ensure permanent
changes behave correctly in multi-month datasets.
"""

import json
import os
from pathlib import Path
from typing import Dict, Any, Optional, Tuple
from datetime import datetime

from ..log_context import get_logger

logger = get_logger(__name__)

# State file location in ~/.czbg directory
STATE_FILE_PATH = Path.home() / ".czbg" / "telemetry_step_state.json"


class StepStateManager:
    """Manages persistent step change state across month boundaries."""
    
    def __init__(self, job_uuid: str):
        """
        Initialize step state manager for a specific generate job.
        
        Args:
            job_uuid: Unique identifier for the current generate job
        """
        self.job_uuid = job_uuid
        self._state_cache = None
        
    def load_state(self) -> Dict[str, Dict[str, Any]]:
        """
        Load step change state from disk.
        
        Returns:
            Dictionary mapping resource_id to step change information
        """
        if self._state_cache is not None:
            return self._state_cache
            
        try:
            if STATE_FILE_PATH.exists():
                with open(STATE_FILE_PATH, 'r') as f:
                    self._state_cache = json.load(f)
                logger.debug(f"Loaded step state for {len(self._state_cache)} resources")
            else:
                self._state_cache = {}
                logger.debug("No existing step state file found")
                
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"Corrupted step state file, recreating: {e}")
            self._state_cache = {}
            self._save_state()
            
        return self._state_cache
    
    def get_resource_state(self, resource_id: str) -> Dict[str, Any]:
        """
        Get step change state for a specific resource.
        
        Args:
            resource_id: Unique identifier for the resource
            
        Returns:
            Dictionary with step change information for the resource
        """
        state = self.load_state()
        return state.get(resource_id, {})
    
    def save_step_baseline(
        self,
        resource_id: str,
        baseline_cost: float,
        step_type: str,
        absolute_increment: float
    ) -> None:
        """
        Save baseline cost and absolute increment for staircase step changes.
        
        Args:
            resource_id: Unique identifier for the resource
            baseline_cost: Original cost before any step changes
            step_type: Type of step change ("step_up" or "step_down")
            absolute_increment: Fixed dollar increment per step
        """
        state = self.load_state()
        
        if resource_id not in state:
            state[resource_id] = {
                "step_up_history": [],
                "step_down_history": [],
                "baseline_cost": None,
                "step_up_absolute_increment": None,
                "step_down_absolute_increment": None,
                "job_uuid": self.job_uuid
            }
        
        # Save baseline and increment
        state[resource_id]["baseline_cost"] = baseline_cost
        state[resource_id][f"{step_type}_absolute_increment"] = absolute_increment
        state[resource_id]["job_uuid"] = self.job_uuid
        
        self._state_cache = state
        self._save_state()
        
        logger.info(f"Saved {step_type} baseline for {resource_id}: ${baseline_cost:.2f}, increment: ${absolute_increment:.2f}")

    def save_step_change(
        self, 
        resource_id: str, 
        step_type: str, 
        step_date: str, 
        step_number: int
    ) -> None:
        """
        Save a step change occurrence for staircase progression.
        
        Args:
            resource_id: Unique identifier for the resource
            step_type: Type of step change ("step_up" or "step_down")
            step_date: Date when step change occurs (YYYY-MM-DD format)
            step_number: Step number (1, 2, 3, etc.)
        """
        state = self.load_state()
        
        if resource_id not in state:
            state[resource_id] = {
                "step_up_history": [],
                "step_down_history": [],
                "baseline_cost": None,
                "step_up_absolute_increment": None,
                "step_down_absolute_increment": None,
                "job_uuid": self.job_uuid
            }
        
        # Add to step history
        step_entry = {"date": step_date, "step_number": step_number}
        state[resource_id][f"{step_type}_history"].append(step_entry)
        state[resource_id]["job_uuid"] = self.job_uuid
        
        self._state_cache = state
        self._save_state()
        
        logger.info(f"Saved {step_type} step {step_number} for {resource_id}: {step_date}")
    
    def has_step_change(self, resource_id: str, step_type: str) -> bool:
        """
        Check if a resource has existing step changes of the given type.
        
        Args:
            resource_id: Unique identifier for the resource
            step_type: Type of step change ("step_up" or "step_down")
            
        Returns:
            True if the resource has any step changes of the specified type
        """
        resource_state = self.get_resource_state(resource_id)
        step_history = resource_state.get(f"{step_type}_history", [])
        return len(step_history) > 0
    
    def get_baseline_cost(self, resource_id: str) -> Optional[float]:
        """
        Get the baseline cost for a resource.
        
        Args:
            resource_id: Unique identifier for the resource
            
        Returns:
            Baseline cost if available, None otherwise
        """
        resource_state = self.get_resource_state(resource_id)
        return resource_state.get("baseline_cost")
    
    def get_absolute_increment(self, resource_id: str, step_type: str) -> Optional[float]:
        """
        Get the absolute increment for step changes.
        
        Args:
            resource_id: Unique identifier for the resource
            step_type: Type of step change ("step_up" or "step_down")
            
        Returns:
            Absolute increment if available, None otherwise
        """
        resource_state = self.get_resource_state(resource_id)
        return resource_state.get(f"{step_type}_absolute_increment")
    
    def should_create_new_step(self, resource_id: str, step_type: str, step_date: str) -> bool:
        """
        Check if a new step change should be created for the given date.
        
        Args:
            resource_id: Unique identifier for the resource
            step_type: Type of step change ("step_up" or "step_down")
            step_date: Date when step change would occur (YYYY-MM-DD format)
            
        Returns:
            True if no step change exists for this 90-day period
        """
        resource_state = self.get_resource_state(resource_id)
        step_history = resource_state.get(f"{step_type}_history", [])
        
        # Convert step_date to datetime for comparison
        from datetime import datetime, timedelta
        target_date = datetime.strptime(step_date, "%Y-%m-%d")
        
        # Check if any existing step falls within the same 90-day period
        for step_entry in step_history:
            existing_date = datetime.strptime(step_entry["date"], "%Y-%m-%d")
            # If within 90 days of existing step, don't create new one
            if abs((target_date - existing_date).days) < 90:
                return False
                
        return True
    
    def cleanup_state(self) -> None:
        """
        Clean up the state file when generate job completes.
        
        Removes the state file to ensure fresh state for next job.
        """
        try:
            if STATE_FILE_PATH.exists():
                STATE_FILE_PATH.unlink()
                logger.info("Cleaned up anomaly step state file")
            else:
                logger.debug("No step state file to clean up")
                
        except OSError as e:
            logger.warning(f"Failed to clean up step state file: {e}")
    
    def _save_state(self) -> None:
        """Save current state to disk."""
        try:
            # Ensure directory exists
            STATE_FILE_PATH.parent.mkdir(parents=True, exist_ok=True)
            
            # Write state file atomically
            temp_path = STATE_FILE_PATH.with_suffix('.tmp')
            with open(temp_path, 'w') as f:
                json.dump(self._state_cache, f, indent=2)
            
            temp_path.replace(STATE_FILE_PATH)
            logger.debug("Saved step state to disk")
            
        except OSError as e:
            logger.error(f"Failed to save step state: {e}")
            raise


# Global instance - will be initialized by generate engine
step_state_manager: Optional[StepStateManager] = None


def initialize_step_state_manager(job_uuid: str) -> None:
    """Initialize the global step state manager for the current job."""
    global step_state_manager
    step_state_manager = StepStateManager(job_uuid)


def cleanup_step_state() -> None:
    """Clean up step state when generate job completes."""
    global step_state_manager
    if step_state_manager:
        step_state_manager.cleanup_state()
        step_state_manager = None


def get_step_state_manager() -> Optional[StepStateManager]:
    """Get the current step state manager instance."""
    return step_state_manager