"""
Directory management utilities for PrecipGen organized output structure.

This module provides utilities for creating and managing the organized directory
structure used by PrecipGen workflows: ./precipgen_output/{station_id}/
"""

import json
import shutil
from pathlib import Path
from typing import List, Optional, Dict, Any
from dataclasses import dataclass
from datetime import datetime

from precipgen.utils import get_logger


@dataclass
class StationDirectory:
    """Information about a station's output directory structure."""
    station_id: str
    directory_path: Path
    data_files: List[str]
    analysis_files: List[str]
    simulation_files: List[str]
    validation_files: List[str]
    total_size_mb: float
    last_modified: datetime


class DirectoryManager:
    """
    Manages organized directory structure for PrecipGen workflow outputs.
    
    Creates and maintains the standard directory structure:
    ./precipgen_output/{station_id}/
    ├── data/
    ├── analysis/
    ├── simulation/
    └── validation/
    """
    
    def __init__(self, base_output_dir: Optional[Path] = None):
        """
        Initialize directory manager.
        
        Args:
            base_output_dir: Base directory for all outputs (default: ./precipgen_output)
        """
        self.base_output_dir = Path(base_output_dir) if base_output_dir else Path("./precipgen_output")
        self.logger = get_logger(__name__)
        
        # Standard subdirectory names
        self.subdirectories = ['data', 'analysis', 'simulation', 'validation']
    
    def create_station_directory(self, station_id: str, base_dir: Optional[Path] = None) -> Path:
        """
        Create organized directory structure for a station.
        
        Args:
            station_id: Station identifier
            base_dir: Override base directory (optional)
            
        Returns:
            Path to the station directory
        """
        if base_dir:
            station_dir = base_dir / station_id
        else:
            station_dir = self.base_output_dir / station_id
        
        # Create station directory
        station_dir.mkdir(parents=True, exist_ok=True)
        
        # Create standard subdirectories
        for subdir in self.subdirectories:
            (station_dir / subdir).mkdir(exist_ok=True)
        
        self.logger.debug(f"Created station directory structure: {station_dir}")
        return station_dir
    
    def get_subdirectory(self, station_id: str, workflow: str, base_dir: Optional[Path] = None) -> Path:
        """
        Get path to a specific workflow subdirectory for a station.
        
        Args:
            station_id: Station identifier
            workflow: Workflow name (data, analysis, simulation, validation)
            base_dir: Override base directory (optional)
            
        Returns:
            Path to the workflow subdirectory
        """
        if workflow not in self.subdirectories:
            raise ValueError(f"Unknown workflow: {workflow}. Must be one of {self.subdirectories}")
        
        station_dir = self.create_station_directory(station_id, base_dir)
        return station_dir / workflow
    
    def list_station_directories(self, base_dir: Optional[Path] = None) -> List[StationDirectory]:
        """
        List all station directories with their information.
        
        Args:
            base_dir: Override base directory (optional)
            
        Returns:
            List of StationDirectory objects
        """
        search_dir = base_dir if base_dir else self.base_output_dir
        
        if not search_dir.exists():
            return []
        
        station_dirs = []
        
        for station_path in search_dir.iterdir():
            if station_path.is_dir() and not station_path.name.startswith('.'):
                try:
                    station_info = self._analyze_station_directory(station_path)
                    station_dirs.append(station_info)
                except Exception as e:
                    self.logger.warning(f"Error analyzing station directory {station_path}: {e}")
        
        return sorted(station_dirs, key=lambda x: x.station_id)
    
    def clean_station_data(self, station_id: str, confirm: bool = True, base_dir: Optional[Path] = None) -> bool:
        """
        Remove all files for a specific station.
        
        Args:
            station_id: Station identifier
            confirm: Whether to require confirmation (default: True)
            base_dir: Override base directory (optional)
            
        Returns:
            True if deletion was successful, False otherwise
        """
        search_dir = base_dir if base_dir else self.base_output_dir
        station_dir = search_dir / station_id
        
        if not station_dir.exists():
            self.logger.warning(f"Station directory does not exist: {station_dir}")
            return False
        
        if confirm:
            # In CLI context, this would prompt user
            # For now, we'll just log the action
            self.logger.info(f"Would delete station directory: {station_dir}")
            return True
        
        try:
            shutil.rmtree(station_dir)
            self.logger.info(f"Deleted station directory: {station_dir}")
            return True
        except Exception as e:
            self.logger.error(f"Failed to delete station directory {station_dir}: {e}")
            return False
    
    def get_station_summary(self, station_id: str, base_dir: Optional[Path] = None) -> Dict[str, Any]:
        """
        Get summary information for a station.
        
        Args:
            station_id: Station identifier
            base_dir: Override base directory (optional)
            
        Returns:
            Dictionary with station summary information
        """
        search_dir = base_dir if base_dir else self.base_output_dir
        station_dir = search_dir / station_id
        
        if not station_dir.exists():
            return {
                'station_id': station_id,
                'exists': False,
                'error': 'Station directory not found'
            }
        
        try:
            station_info = self._analyze_station_directory(station_dir)
            
            return {
                'station_id': station_id,
                'exists': True,
                'directory_path': str(station_info.directory_path),
                'total_size_mb': station_info.total_size_mb,
                'last_modified': station_info.last_modified.isoformat(),
                'file_counts': {
                    'data': len(station_info.data_files),
                    'analysis': len(station_info.analysis_files),
                    'simulation': len(station_info.simulation_files),
                    'validation': len(station_info.validation_files)
                },
                'workflows_completed': self._get_completed_workflows(station_info)
            }
        
        except Exception as e:
            return {
                'station_id': station_id,
                'exists': True,
                'error': str(e)
            }
    
    def create_timestamped_subdirectory(self, station_id: str, workflow: str, base_dir: Optional[Path] = None) -> Path:
        """
        Create a timestamped subdirectory within a workflow directory.
        
        Args:
            station_id: Station identifier
            workflow: Workflow name
            base_dir: Override base directory (optional)
            
        Returns:
            Path to the timestamped subdirectory
        """
        workflow_dir = self.get_subdirectory(station_id, workflow, base_dir)
        
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        timestamped_dir = workflow_dir / timestamp
        timestamped_dir.mkdir(parents=True, exist_ok=True)
        
        return timestamped_dir
    
    def _analyze_station_directory(self, station_path: Path) -> StationDirectory:
        """Analyze a station directory and return information about it."""
        station_id = station_path.name
        
        # Get file lists for each subdirectory
        data_files = self._get_files_in_subdir(station_path / 'data')
        analysis_files = self._get_files_in_subdir(station_path / 'analysis')
        simulation_files = self._get_files_in_subdir(station_path / 'simulation')
        validation_files = self._get_files_in_subdir(station_path / 'validation')
        
        # Calculate total size
        total_size = self._calculate_directory_size(station_path)
        
        # Get last modified time
        last_modified = datetime.fromtimestamp(station_path.stat().st_mtime)
        
        return StationDirectory(
            station_id=station_id,
            directory_path=station_path,
            data_files=data_files,
            analysis_files=analysis_files,
            simulation_files=simulation_files,
            validation_files=validation_files,
            total_size_mb=total_size / (1024 * 1024),
            last_modified=last_modified
        )
    
    def _get_files_in_subdir(self, subdir_path: Path) -> List[str]:
        """Get list of files in a subdirectory."""
        if not subdir_path.exists():
            return []
        
        files = []
        for item in subdir_path.rglob('*'):
            if item.is_file():
                # Get relative path from subdirectory
                rel_path = item.relative_to(subdir_path)
                files.append(str(rel_path))
        
        return sorted(files)
    
    def _calculate_directory_size(self, directory: Path) -> int:
        """Calculate total size of directory in bytes."""
        total_size = 0
        
        for item in directory.rglob('*'):
            if item.is_file():
                try:
                    total_size += item.stat().st_size
                except (OSError, FileNotFoundError):
                    # Skip files that can't be accessed
                    pass
        
        return total_size
    
    def get_station_info(self, station_id: str, base_dir: Optional[Path] = None) -> Optional[StationDirectory]:
        """
        Get detailed information about a station directory.
        
        Args:
            station_id: Station identifier
            base_dir: Override base directory (optional)
            
        Returns:
            StationDirectory object if station exists, None otherwise
        """
        search_dir = base_dir if base_dir else self.base_output_dir
        station_dir = search_dir / station_id
        
        if not station_dir.exists():
            return None
        
        try:
            return self._analyze_station_directory(station_dir)
        except Exception as e:
            self.logger.warning(f"Error analyzing station directory {station_dir}: {e}")
            return None
    
    def _get_completed_workflows(self, station_info: StationDirectory) -> List[str]:
        """Determine which workflows have been completed based on file presence."""
        completed = []
        
        # Check for key files that indicate workflow completion
        if any('precipitation.csv' in f for f in station_info.data_files):
            completed.append('data_acquisition')
        
        if any('params.json' in f for f in station_info.analysis_files):
            completed.append('analysis')
        
        if any('synthetic.csv' in f for f in station_info.simulation_files):
            completed.append('simulation')
        
        if any('validation_report.json' in f for f in station_info.validation_files):
            completed.append('validation')
        
        return completed