"""
Infrastructure Pipeline for orchestrating discovery and generation stages.
"""

import json
import os
import glob
from pathlib import Path
from typing import Dict, Any, List
import click

from .features.infra.infra_feature import InfraFeature
from .models.context import (
    InfraPipelineContext, 
    DiscoveryStageContext, 
    GenerationStageContext
)
from .models.cloudformation import CloudFormationTemplate

# Import features to ensure they are registered
from . import features


class InfraPipeline:
    """Main pipeline orchestrating discovery and generation stages."""
    
    def __init__(self):
        self.pipeline_context = InfraPipelineContext()

    def process_directory(self, directory_path: str) -> str:
        """
        Process all Python files in a directory through the pipeline.
        
        Args:
            directory_path: Path to the directory to process
            
        Returns:
            JSON-formatted CloudFormation templates
        """
        stack_templates = self.process_directory_to_stacks(directory_path)
        
        # If only one stack, return it directly
        if len(stack_templates) == 1:
            return list(stack_templates.values())[0]
        
        # If multiple stacks, return them as a JSON object with stack names as keys
        return json.dumps(stack_templates, indent=2)

    def process_directory_to_stacks(self, directory_path: str) -> Dict[str, str]:
        """
        Process all Python files in a directory and return separate templates for each stack.
        
        Args:
            directory_path: Path to the directory to process
            
        Returns:
            Dictionary mapping stack names to JSON-formatted CloudFormation templates
        """
        # Get all registered features
        feature_instances = [feature_class() for feature_class in InfraFeature._FEATURES.values()]
        
        # Process each Python file
        python_files = glob.glob(os.path.join(directory_path, '**', '*.py'), recursive=True)
        
        for file_path in python_files:
            self.process_file(Path(file_path))
        
        # Build CloudFormation templates for each stack
        return self._build_stack_templates()

    def process_file(self, file_path: Path):
        """Process a single file through the pipeline stages"""
        
        # Stage 1: Discovery
        discovery_context = DiscoveryStageContext(file_path=file_path)
        
        # Store discovered configs by feature class and config type
        file_feature_configs = {}
        
        for feature_class in InfraFeature._FEATURES.values():
            feature = feature_class()
            discovered_configs = feature.discover(discovery_context)
            
            if discovered_configs:  # Only store if configs were found
                # Group configs by their type name (string)
                configs_by_type_name = {}
                for config in discovered_configs:
                    config_type_name = type(config).__name__
                    if config_type_name not in configs_by_type_name:
                        configs_by_type_name[config_type_name] = []
                    configs_by_type_name[config_type_name].append(config)
                
                file_feature_configs[feature_class.__name__] = configs_by_type_name
                
                # Also store in pipeline context for global access
                self.pipeline_context.add_discovered_configs(feature_class.__name__, discovered_configs)
        
        # Stage 2: Generation
        generation_context = GenerationStageContext(
            file_path=file_path,
            feature_configs=file_feature_configs
        )
        
        # All features generate CloudFormation from their discovered configs
        for feature_class in InfraFeature._FEATURES.values():
            feature = feature_class()
            generated_stacks = feature.generate(generation_context)
            
            # Pipeline manages merging generated stacks
            self._merge_stacks(generated_stacks)

    def _merge_stacks(self, generated_stacks: Dict[str, CloudFormationTemplate]) -> None:
        """Merge generated stacks into pipeline context"""
        for stack_name, stack_template in generated_stacks.items():
            if stack_name not in self.pipeline_context.stacks:
                self.pipeline_context.stacks[stack_name] = {"Resources": {}, "Outputs": {}}
            
            # Merge resources
            for resource_id, resource in stack_template.resources.items():
                self.pipeline_context.stacks[stack_name]["Resources"][resource_id] = {
                    "Type": resource.type,
                    "Properties": resource.model_dump()["properties"]
                }
            
            # Merge outputs
            for output_id, output in stack_template.outputs.items():
                self.pipeline_context.stacks[stack_name]["Outputs"][output_id] = {
                    "Description": output.description,
                    "Value": output.value.model_dump(),
                    "Export": {"Name": output.export.name}
                }

    def _build_stack_templates(self) -> Dict[str, str]:
        """Build CloudFormation templates for each stack"""
        stack_templates = {}
        
        for stack_name, stack_data in self.pipeline_context.stacks.items():
            cf_template = {
                "AWSTemplateFormatVersion": "2010-09-09",
                "Description": f"CloudFormation template generated by iicy pipeline for stack {stack_name}",
                "Resources": stack_data.get("Resources", {}),
                "Outputs": stack_data.get("Outputs", {})
            }
            stack_templates[stack_name] = json.dumps(cf_template, indent=2)
        
        # If no stacks found, return default empty stack
        if not stack_templates:
            default_stack = {
                "AWSTemplateFormatVersion": "2010-09-09",
                "Description": "Empty CloudFormation template - no resources found",
                "Resources": {},
                "Outputs": {}
            }
            stack_templates["IICYStack"] = json.dumps(default_stack, indent=2)
        
        return stack_templates

    # Main public API method
    def generate(self, path: str) -> Dict[str, str]:
        """
        Main entry point for generating CloudFormation templates from a path.
        Automatically handles both files and directories.
        
        Args:
            path: Path to a file or directory to analyze
            
        Returns:
            Dictionary mapping stack names to JSON-formatted CloudFormation templates
        """
        path_obj = Path(path)
        
        if path_obj.is_file():
            # Single file
            self.process_file(path_obj)
        elif path_obj.is_dir():
            # Directory - process all Python files
            python_files = glob.glob(os.path.join(path, '**', '*.py'), recursive=True)
            for file_path in python_files:
                self.process_file(Path(file_path))
        else:
            raise ValueError(f"Path does not exist: {path}")
        
        return self._build_stack_templates()

    def get_stack_names(self, path: str) -> List[str]:
        """
        Get stack names that would be generated from analyzing a path.
        Automatically handles both files and directories.
        
        Args:
            path: Path to a file or directory to analyze
            
        Returns:
            List of stack names that would be generated
        """
        stack_templates = self.generate(path)
        return list(stack_templates.keys())

    # Public API methods (formerly in generator.py)
    
    def analyze_directory(self, directory_path: str) -> str:
        """
        Analyzes all Python files in a directory using the pipeline.
        
        Args:
            directory_path: Path to the directory to analyze
            
        Returns:
            JSON-formatted CloudFormation templates (one per stack)
        """
        stack_templates = self.process_directory_to_stacks(directory_path)
        
        # If only one stack, return it directly
        if len(stack_templates) == 1:
            return list(stack_templates.values())[0]
        
        # If multiple stacks, return them as a JSON object with stack names as keys
        return json.dumps(stack_templates, indent=2)

    def analyze_file(self, file_path: str) -> str:
        """
        Analyzes a single Python file using the pipeline.
        
        Args:
            file_path: Path to the Python file to analyze
            
        Returns:
            JSON-formatted CloudFormation template
        """
        stack_templates = self.analyze_file_to_stacks(file_path)
        
        # If only one stack, return it directly
        if len(stack_templates) == 1:
            return list(stack_templates.values())[0]
        
        # If multiple stacks, return them as a JSON object with stack names as keys
        return json.dumps(stack_templates, indent=2)

    def analyze_file_to_stacks(self, file_path: str) -> Dict[str, str]:
        """
        Analyzes a single Python file and returns separate templates for each stack.
        
        Args:
            file_path: Path to the Python file to analyze
            
        Returns:
            Dictionary mapping stack names to JSON-formatted CloudFormation templates
        """
        self.process_file(Path(file_path))
        return self._build_stack_templates()

    def get_stack_names_from_path(self, path: str) -> List[str]:
        """
        Determine what stack names would be generated from analyzing a given path.
        
        Args:
            path: File or directory path to analyze
            
        Returns:
            List of stack names that would be generated
        """
        if click.format_filename(path).endswith('.py'):
            stack_templates = self.analyze_file_to_stacks(path)
        else:
            stack_templates = self.process_directory_to_stacks(path)
        
        return list(stack_templates.keys())
