"""
Directory management for organized workflow testing output.

This module provides the DirectoryManager class for creating and managing
station-based directory structures that organize all workflow outputs
according to the PrecipGen testing framework requirements.
"""

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

from precipgen.utils import get_logger


@dataclass
class StationDirectory:
    """Information about a station's output directory."""
    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 workflow testing outputs.
    
    Creates and maintains station-based directory structure:
    ./precipgen_output/{station_id}/
    ├── data/
    ├── analysis/{timestamp}/
    ├── simulation/{timestamp}/
    └── validation/{timestamp}/
    """
    
    def __init__(self, base_dir: str = "./precipgen_output"):
        """
        Initialize DirectoryManager with base output directory.
        
        Args:
            base_dir: Base directory for all workflow outputs
        """
        self.base_dir = Path(base_dir)
        self.logger = get_logger("workflow_testing.directory_manager")
        
        # Ensure base directory exists
        self.base_dir.mkdir(parents=True, exist_ok=True)
        self.logger.info(f"DirectoryManager initialized with base_dir: {self.base_dir}")
    
    def create_station_directory(self, station_id: str, 
                               base_dir: Optional[str] = None) -> Path:
        """
        Create or get station directory with organized subdirectory structure.
        
        Args:
            station_id: GHCN station identifier
            base_dir: Optional custom base directory (overrides default)
            
        Returns:
            Path to station directory
            
        Raises:
            OSError: If directory creation fails
        """
        if base_dir:
            station_base = Path(base_dir) / station_id
        else:
            station_base = self.base_dir / station_id
        
        # Create station directory and subdirectories
        subdirs = ["data", "analysis", "simulation", "validation"]
        
        try:
            station_base.mkdir(parents=True, exist_ok=True)
            
            for subdir in subdirs:
                (station_base / subdir).mkdir(exist_ok=True)
            
            self.logger.info(f"Created station directory structure for {station_id}: {station_base}")
            return station_base
            
        except OSError as e:
            self.logger.error(f"Failed to create station directory for {station_id}: {e}")
            raise
    
    def get_subdirectory(self, station_id: str, workflow: str, 
                        timestamp: Optional[str] = None,
                        base_dir: Optional[str] = None) -> Path:
        """
        Get workflow-specific subdirectory, creating timestamped subdirs as needed.
        
        Args:
            station_id: GHCN station identifier
            workflow: Workflow name (data, analysis, simulation, validation)
            timestamp: Optional timestamp for versioned outputs
            base_dir: Optional custom base directory
            
        Returns:
            Path to workflow subdirectory
            
        Raises:
            ValueError: If workflow name is invalid
        """
        valid_workflows = ["data", "analysis", "simulation", "validation"]
        if workflow not in valid_workflows:
            raise ValueError(f"Invalid workflow '{workflow}'. Must be one of: {valid_workflows}")
        
        # Get station directory
        station_dir = self.create_station_directory(station_id, base_dir)
        
        # For data workflow, no timestamp subdirectory
        if workflow == "data":
            return station_dir / workflow
        
        # For other workflows, create timestamped subdirectory if specified
        workflow_dir = station_dir / workflow
        
        if timestamp:
            timestamped_dir = workflow_dir / timestamp
            timestamped_dir.mkdir(parents=True, exist_ok=True)
            return timestamped_dir
        
        return workflow_dir
    
    def create_timestamped_subdirectory(self, station_id: str, workflow: str,
                                      base_dir: Optional[str] = None) -> Path:
        """
        Create timestamped subdirectory for workflow outputs.
        
        Args:
            station_id: GHCN station identifier
            workflow: Workflow name (analysis, simulation, validation)
            base_dir: Optional custom base directory
            
        Returns:
            Path to timestamped subdirectory
        """
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return self.get_subdirectory(station_id, workflow, timestamp, base_dir)
    
    def list_station_directories(self, base_dir: Optional[str] = None) -> List[StationDirectory]:
        """
        List all station directories with file information.
        
        Args:
            base_dir: Optional custom base directory
            
        Returns:
            List of StationDirectory objects with metadata
        """
        search_dir = Path(base_dir) if base_dir else self.base_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"Failed to analyze station directory {station_path}: {e}")
        
        return sorted(station_dirs, key=lambda x: x.station_id)
    
    def _analyze_station_directory(self, station_path: Path) -> StationDirectory:
        """
        Analyze a station directory and collect file information.
        
        Args:
            station_path: Path to station directory
            
        Returns:
            StationDirectory object with metadata
        """
        station_id = station_path.name
        
        # Collect files by workflow type
        data_files = self._collect_files(station_path / "data")
        analysis_files = self._collect_files(station_path / "analysis")
        simulation_files = self._collect_files(station_path / "simulation")
        validation_files = self._collect_files(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,
            last_modified=last_modified
        )
    
    def _collect_files(self, directory: Path) -> List[str]:
        """
        Recursively collect all files in a directory.
        
        Args:
            directory: Directory to scan
            
        Returns:
            List of relative file paths
        """
        if not directory.exists():
            return []
        
        files = []
        for item in directory.rglob("*"):
            if item.is_file():
                relative_path = item.relative_to(directory)
                files.append(str(relative_path))
        
        return sorted(files)
    
    def _calculate_directory_size(self, directory: Path) -> float:
        """
        Calculate total size of directory in MB.
        
        Args:
            directory: Directory to measure
            
        Returns:
            Size in megabytes
        """
        total_size = 0
        
        try:
            for item in directory.rglob("*"):
                if item.is_file():
                    total_size += item.stat().st_size
        except (OSError, PermissionError):
            # Handle permission errors gracefully
            pass
        
        return total_size / (1024 * 1024)  # Convert to MB
    
    def clean_station_data(self, station_id: str, confirm: bool = True,
                          base_dir: Optional[str] = None) -> bool:
        """
        Remove all files for a specific station with optional confirmation.
        
        Args:
            station_id: GHCN station identifier
            confirm: Whether to require confirmation (for CLI usage)
            base_dir: Optional custom base directory
            
        Returns:
            True if deletion was successful, False otherwise
        """
        search_dir = Path(base_dir) if base_dir else self.base_dir
        station_path = search_dir / station_id
        
        if not station_path.exists():
            self.logger.warning(f"Station directory does not exist: {station_path}")
            return False
        
        if confirm:
            # In actual CLI usage, this would prompt the user
            # For testing, we'll log the action
            self.logger.info(f"Would prompt for confirmation to delete: {station_path}")
        
        try:
            shutil.rmtree(station_path)
            self.logger.info(f"Successfully deleted station directory: {station_path}")
            return True
            
        except OSError as e:
            self.logger.error(f"Failed to delete station directory {station_path}: {e}")
            return False
    
    def get_station_info(self, station_id: str, 
                        base_dir: Optional[str] = None) -> Optional[StationDirectory]:
        """
        Get detailed information about a specific station directory.
        
        Args:
            station_id: GHCN station identifier
            base_dir: Optional custom base directory
            
        Returns:
            StationDirectory object or None if not found
        """
        search_dir = Path(base_dir) if base_dir else self.base_dir
        station_path = search_dir / station_id
        
        if not station_path.exists():
            return None
        
        try:
            return self._analyze_station_directory(station_path)
        except Exception as e:
            self.logger.error(f"Failed to analyze station directory {station_path}: {e}")
            return None
    
    def ensure_workflow_directory(self, station_id: str, workflow: str,
                                 base_dir: Optional[str] = None) -> Path:
        """
        Ensure workflow directory exists and return its path.
        
        Args:
            station_id: GHCN station identifier
            workflow: Workflow name
            base_dir: Optional custom base directory
            
        Returns:
            Path to workflow directory
        """
        workflow_dir = self.get_subdirectory(station_id, workflow, base_dir=base_dir)
        workflow_dir.mkdir(parents=True, exist_ok=True)
        return workflow_dir
    
    def get_output_summary(self, base_dir: Optional[str] = None) -> Dict[str, Any]:
        """
        Get summary statistics for all station outputs.
        
        Args:
            base_dir: Optional custom base directory
            
        Returns:
            Dictionary with summary statistics
        """
        station_dirs = self.list_station_directories(base_dir)
        
        if not station_dirs:
            return {
                "total_stations": 0,
                "total_size_mb": 0.0,
                "stations": []
            }
        
        total_size = sum(station.total_size_mb for station in station_dirs)
        
        return {
            "total_stations": len(station_dirs),
            "total_size_mb": total_size,
            "stations": [
                {
                    "station_id": station.station_id,
                    "size_mb": station.total_size_mb,
                    "file_counts": {
                        "data": len(station.data_files),
                        "analysis": len(station.analysis_files),
                        "simulation": len(station.simulation_files),
                        "validation": len(station.validation_files)
                    },
                    "last_modified": station.last_modified.isoformat()
                }
                for station in station_dirs
            ]
        }
