"""
Multi-Series Setup Utility

This utility helps users set up the proper folder structure and configuration
for multi-series dongle inference.

Features:
- Create recommended folder structure
- Validate existing folder structure  
- Generate example configuration files
- Provide setup guidance and troubleshooting

Usage:
    python utils/multi_series_setup.py create-structure --path "C:/MyAssets"
    python utils/multi_series_setup.py validate --path "C:/MyAssets"
    python utils/multi_series_setup.py help
"""

import os
import sys
import argparse
from typing import List, Tuple, Dict
import json


class MultiSeriesSetup:
    """Utility class for multi-series setup operations"""
    
    SUPPORTED_SERIES = ['520', '720', '630', '730', '540']
    REQUIRED_FW_FILES = ['fw_scpu.bin', 'fw_ncpu.bin']
    
    @staticmethod
    def create_folder_structure(base_path: str, series_list: List[str] = None) -> bool:
        """
        Create the recommended folder structure for multi-series assets
        
        Args:
            base_path: Root path where assets folder should be created
            series_list: List of series to create folders for
            
        Returns:
            bool: Success status
        """
        if series_list is None:
            series_list = MultiSeriesSetup.SUPPORTED_SERIES
        
        try:
            assets_path = os.path.join(base_path, 'Assets')
            firmware_path = os.path.join(assets_path, 'Firmware')
            models_path = os.path.join(assets_path, 'Models')
            
            # Create main directories
            os.makedirs(firmware_path, exist_ok=True)
            os.makedirs(models_path, exist_ok=True)
            
            print(f"✓ Created main directories at: {assets_path}")
            
            # Create series-specific directories
            created_series = []
            for series in series_list:
                series_name = f'KL{series}'
                fw_dir = os.path.join(firmware_path, series_name)
                model_dir = os.path.join(models_path, series_name)
                
                os.makedirs(fw_dir, exist_ok=True)
                os.makedirs(model_dir, exist_ok=True)
                created_series.append(series_name)
                
            print(f"✓ Created series directories: {', '.join(created_series)}")
            
            # Create README file explaining the structure
            readme_content = MultiSeriesSetup._generate_readme_content()
            readme_path = os.path.join(assets_path, 'README.md')
            with open(readme_path, 'w') as f:
                f.write(readme_content)
            
            print(f"✓ Created README file: {readme_path}")
            
            # Create example configuration
            config_example = MultiSeriesSetup._generate_example_config(assets_path, series_list)
            config_path = os.path.join(assets_path, 'example_config.json')
            with open(config_path, 'w') as f:
                json.dump(config_example, f, indent=2)
            
            print(f"✓ Created example configuration: {config_path}")
            
            print(f"\n🎉 Multi-series folder structure created successfully!")
            print(f"📁 Assets folder: {assets_path}")
            print("\n📋 Next steps:")
            print("1. Copy your firmware files to the appropriate Firmware/KLxxx/ folders")
            print("2. Copy your model files to the appropriate Models/KLxxx/ folders") 
            print("3. Configure your model node to use this Assets folder")
            print("4. Enable multi-series mode and select desired series")
            
            return True
            
        except Exception as e:
            print(f"❌ Error creating folder structure: {e}")
            return False
    
    @staticmethod
    def validate_folder_structure(assets_path: str) -> Tuple[bool, List[str]]:
        """
        Validate an existing folder structure for multi-series configuration
        
        Args:
            assets_path: Path to the assets folder
            
        Returns:
            Tuple of (is_valid, list_of_issues)
        """
        issues = []
        
        # Check if assets folder exists
        if not os.path.exists(assets_path):
            issues.append(f"Assets folder does not exist: {assets_path}")
            return False, issues
        
        # Check for main folders
        firmware_path = os.path.join(assets_path, 'Firmware')
        models_path = os.path.join(assets_path, 'Models')
        
        if not os.path.exists(firmware_path):
            issues.append(f"Firmware folder missing: {firmware_path}")
        
        if not os.path.exists(models_path):
            issues.append(f"Models folder missing: {models_path}")
        
        if issues:
            return False, issues
        
        # Check series folders and contents
        found_series = []
        
        for item in os.listdir(firmware_path):
            if item.startswith('KL') and os.path.isdir(os.path.join(firmware_path, item)):
                series_name = item
                series_fw_path = os.path.join(firmware_path, series_name)
                series_model_path = os.path.join(models_path, series_name)
                
                # Check if corresponding model folder exists
                if not os.path.exists(series_model_path):
                    issues.append(f"Models folder missing for {series_name}: {series_model_path}")
                    continue
                
                # Check firmware files
                fw_issues = []
                for fw_file in MultiSeriesSetup.REQUIRED_FW_FILES:
                    fw_file_path = os.path.join(series_fw_path, fw_file)
                    if not os.path.exists(fw_file_path):
                        fw_issues.append(f"{fw_file} missing")
                
                if fw_issues:
                    issues.append(f"{series_name} firmware issues: {', '.join(fw_issues)}")
                
                # Check for model files
                model_files = [f for f in os.listdir(series_model_path) if f.endswith('.nef')]
                if not model_files:
                    issues.append(f"{series_name} has no .nef model files in {series_model_path}")
                
                if not fw_issues and model_files:
                    found_series.append(series_name)
        
        if not found_series:
            issues.append("No valid series configurations found")
        
        is_valid = len(issues) == 0
        
        # Print validation results
        if is_valid:
            print(f"✅ Validation passed!")
            print(f"📁 Assets folder: {assets_path}")
            print(f"🎯 Valid series found: {', '.join(found_series)}")
        else:
            print(f"❌ Validation failed with {len(issues)} issues:")
            for i, issue in enumerate(issues, 1):
                print(f"   {i}. {issue}")
                
        return is_valid, issues
    
    @staticmethod
    def list_available_series(assets_path: str) -> Dict[str, Dict[str, any]]:
        """
        List all available and configured series in the assets folder
        
        Args:
            assets_path: Path to the assets folder
            
        Returns:
            Dict with series information
        """
        series_info = {}
        
        if not os.path.exists(assets_path):
            return series_info
        
        firmware_path = os.path.join(assets_path, 'Firmware')
        models_path = os.path.join(assets_path, 'Models')
        
        if not os.path.exists(firmware_path) or not os.path.exists(models_path):
            return series_info
        
        for item in os.listdir(firmware_path):
            if item.startswith('KL') and os.path.isdir(os.path.join(firmware_path, item)):
                series_name = item
                series_fw_path = os.path.join(firmware_path, series_name)
                series_model_path = os.path.join(models_path, series_name)
                
                # Check firmware files
                fw_files = {}
                for fw_file in MultiSeriesSetup.REQUIRED_FW_FILES:
                    fw_file_path = os.path.join(series_fw_path, fw_file)
                    fw_files[fw_file] = os.path.exists(fw_file_path)
                
                # Check model files
                model_files = []
                if os.path.exists(series_model_path):
                    model_files = [f for f in os.listdir(series_model_path) if f.endswith('.nef')]
                
                # Determine status
                fw_complete = all(fw_files.values())
                has_models = len(model_files) > 0
                
                if fw_complete and has_models:
                    status = "✅ Ready"
                elif fw_complete:
                    status = "⚠️ Missing models"
                elif has_models:
                    status = "⚠️ Missing firmware"
                else:
                    status = "❌ Incomplete"
                
                series_info[series_name] = {
                    'status': status,
                    'firmware_files': fw_files,
                    'model_files': model_files,
                    'firmware_path': series_fw_path,
                    'models_path': series_model_path
                }
        
        return series_info
    
    @staticmethod
    def _generate_readme_content() -> str:
        """Generate README content for the assets folder"""
        return '''
# Multi-Series Assets Folder Structure

This folder contains firmware and models organized by dongle series for multi-series inference.

## Structure:
```
Assets/
├── Firmware/
│   ├── KL520/
│   │   ├── fw_scpu.bin
│   │   └── fw_ncpu.bin
│   ├── KL720/
│   │   ├── fw_scpu.bin
│   │   └── fw_ncpu.bin
│   └── [other series...]
└── Models/
    ├── KL520/
    │   └── [model.nef files]
    ├── KL720/
    │   └── [model.nef files]
    └── [other series...]
```

## Usage:
1. Place firmware files (fw_scpu.bin, fw_ncpu.bin) in the appropriate series subfolder under Firmware/
2. Place model files (.nef) in the appropriate series subfolder under Models/
3. Configure your model node to use this Assets folder in multi-series mode
4. Select which series to enable in the model node properties

## Supported Series:
- **KL520**: Entry-level performance (3 GOPS)
- **KL720**: Mid-range performance (28 GOPS)
- **KL630**: High performance (400 GOPS)
- **KL730**: Very high performance (1600 GOPS)
- **KL540**: Specialized performance (800 GOPS)

## Performance Benefits:
The multi-series system automatically load balances inference across all enabled series
based on their GOPS capacity for optimal performance. You can expect:

- Higher overall throughput by utilizing multiple dongle types simultaneously
- Automatic load balancing based on dongle capabilities
- Seamless failover if one series becomes unavailable
- Scalable performance as you add more dongles

## Validation:
Run `python utils/multi_series_setup.py validate --path <this_folder>` to validate your configuration.

## Troubleshooting:
- Ensure all firmware files are exactly named `fw_scpu.bin` and `fw_ncpu.bin`
- Model files must have `.nef` extension
- Each series must have both firmware and at least one model file
- Check file permissions and accessibility
        '''.strip()
    
    @staticmethod
    def _generate_example_config(assets_path: str, series_list: List[str]) -> Dict:
        """Generate example configuration for model node"""
        return {
            "model_node_properties": {
                "multi_series_mode": True,
                "assets_folder": assets_path,
                "enabled_series": series_list[:2],  # Enable first two series by default
                "max_queue_size": 100,
                "result_buffer_size": 1000,
                "batch_size": 1
            },
            "expected_performance": {
                "total_gops": sum([
                    {"520": 3, "720": 28, "630": 400, "730": 1600, "540": 800}.get(series, 0)
                    for series in series_list[:2]
                ]),
                "load_balancing": "automatic",
                "expected_fps_improvement": "2-5x compared to single series"
            },
            "notes": [
                "This is an example configuration",
                "Adjust enabled_series based on your available dongles", 
                "Higher queue sizes may improve performance but use more memory",
                "Monitor system resources when using multiple series"
            ]
        }


def main():
    """Main CLI interface"""
    parser = argparse.ArgumentParser(description='Multi-Series Dongle Setup Utility')
    subparsers = parser.add_subparsers(dest='command', help='Available commands')
    
    # Create structure command
    create_parser = subparsers.add_parser('create-structure', help='Create folder structure')
    create_parser.add_argument('--path', required=True, help='Base path for assets folder')
    create_parser.add_argument('--series', nargs='*', default=None, help='Series to set up (default: all)')
    
    # Validate command
    validate_parser = subparsers.add_parser('validate', help='Validate existing structure')
    validate_parser.add_argument('--path', required=True, help='Path to assets folder')
    
    # List command
    list_parser = subparsers.add_parser('list', help='List available series')
    list_parser.add_argument('--path', required=True, help='Path to assets folder')
    
    # Help command
    help_parser = subparsers.add_parser('help', help='Show detailed help')
    
    args = parser.parse_args()
    
    if args.command == 'create-structure':
        series_list = args.series if args.series else None
        MultiSeriesSetup.create_folder_structure(args.path, series_list)
    
    elif args.command == 'validate':
        is_valid, issues = MultiSeriesSetup.validate_folder_structure(args.path)
        sys.exit(0 if is_valid else 1)
    
    elif args.command == 'list':
        series_info = MultiSeriesSetup.list_available_series(args.path)
        
        if not series_info:
            print(f"❌ No series found in {args.path}")
            sys.exit(1)
        
        print(f"📁 Series configuration in {args.path}:\n")
        
        for series_name, info in series_info.items():
            print(f"  {series_name}: {info['status']}")
            print(f"    📁 Firmware: {info['firmware_path']}")
            print(f"    📁 Models: {info['models_path']}")
            
            if info['model_files']:
                print(f"    📄 Model files: {', '.join(info['model_files'])}")
            
            fw_issues = [fw for fw, exists in info['firmware_files'].items() if not exists]
            if fw_issues:
                print(f"    ⚠️  Missing firmware: {', '.join(fw_issues)}")
            
            print()
    
    elif args.command == 'help':
        print("""
Multi-Series Dongle Setup Help
=============================

This utility helps you set up and manage multi-series dongle configurations
for improved inference performance.

Commands:
---------

create-structure --path <path> [--series KL520 KL720 ...]
    Creates the recommended folder structure for multi-series assets.
    
    Example:
        python utils/multi_series_setup.py create-structure --path "C:/MyAssets"
        python utils/multi_series_setup.py create-structure --path "C:/MyAssets" --series 520 720

validate --path <path>
    Validates an existing assets folder structure.
    
    Example:
        python utils/multi_series_setup.py validate --path "C:/MyAssets/Assets"

list --path <path>
    Lists all available series and their status in an assets folder.
    
    Example:
        python utils/multi_series_setup.py list --path "C:/MyAssets/Assets"

Setup Workflow:
--------------
1. Create folder structure: create-structure --path "C:/MyProject"
2. Copy firmware files to Assets/Firmware/KLxxx/ folders
3. Copy model files to Assets/Models/KLxxx/ folders  
4. Validate configuration: validate --path "C:/MyProject/Assets"
5. Configure model node in UI to use Assets folder
6. Enable multi-series mode and select desired series

Performance Benefits:
-------------------
- 2-5x throughput improvement with multiple series
- Automatic load balancing based on dongle GOPS
- Seamless scaling as you add more dongles
- Fault tolerance if some dongles become unavailable

Troubleshooting:
---------------
- Ensure exact firmware file names: fw_scpu.bin, fw_ncpu.bin
- Model files must have .nef extension
- Check file permissions and paths
- Verify dongle connectivity with single-series mode first
- Use validate command to check configuration

For more help, see Assets/README.md after creating the structure.
        """)
    
    else:
        parser.print_help()


if __name__ == '__main__':
    main()