"""
Pipeline context for infrastructure features.
"""

from pathlib import Path
from pydantic import BaseModel, Field
from typing import Dict, Any, TypeVar, Type, List, Optional, TYPE_CHECKING
from .cloudformation import Stack

if TYPE_CHECKING:
    from ..features.infra.infra_feature import InfraFeature

T = TypeVar('T')

class InfraPipelineContext(BaseModel):
    """
    Main pipeline context - exclusively owned by InfraPipeline.
    Contains all discovered configs and generated stacks across all files.
    """
    stacks: Dict[str, Any] = Field(default_factory=dict)  # Raw CloudFormation template data
    resource_mappings: Dict[str, Dict[str, tuple]] = Field(default_factory=dict)
    discovered_configs: Dict[str, List[Any]] = Field(default_factory=dict)  # All discovered configs across all files
    
    def add_discovered_configs(self, feature_name: str, configs: List[Any]) -> None:
        """Add discovered config instances for a feature"""
        if feature_name not in self.discovered_configs:
            self.discovered_configs[feature_name] = []
        self.discovered_configs[feature_name].extend(configs)

class DiscoveryStageContext(BaseModel):
    """Context for discovery stage - only contains what discovery needs"""
    file_path: Path

class FeatureAccessor:
    """Accessor for a specific feature's discovered configs"""
    
    def __init__(self, configs_by_type: Dict[Type, List[Any]]):
        self._configs_by_type = configs_by_type
    
    def resources(self, config_class: Type[T]) -> List[T]:
        """
        Get all discovered config instances of a specific type.
        
        Args:
            config_class: The config class type to retrieve
            
        Returns:
            List of discovered config instances of the specified type
        """
        return self._configs_by_type.get(config_class, [])
    
    def resource(self, config_class: Type[T], name: str) -> Optional[T]:
        """
        Get a specific discovered config instance by name.
        
        Args:
            config_class: The config class type to retrieve
            name: The name/identifier of the specific resource
            
        Returns:
            The specific config instance, or None if not found
        """
        configs = self._configs_by_type.get(config_class, [])
        for config in configs:
            # Try common name fields
            for name_field in ['name', 'logical_id', 'function_name']:
                if hasattr(config, name_field) and getattr(config, name_field) == name:
                    return config
        return None
    
    def has_resource(self, config_class: Type[T], name: str) -> bool:
        """
        Check if a specific resource exists.
        
        Args:
            config_class: The config class type to check
            name: The name/identifier of the resource
            
        Returns:
            True if the resource exists, False otherwise
        """
        return self.resource(config_class, name) is not None
    
    def count_resources(self, config_class: Type[T]) -> int:
        """
        Get the count of discovered config instances of a specific type.
        
        Args:
            config_class: The config class type to count
            
        Returns:
            Number of discovered config instances
        """
        return len(self._configs_by_type.get(config_class, []))

class GenerationStageContext(BaseModel):
    """Context for generation stage with feature-based config access"""
    file_path: Path
    
    # Internal storage - not directly accessible
    feature_configs: Dict[str, Dict[str, List[Any]]] = Field(default_factory=dict)
    
    def feature(self, feature_class) -> FeatureAccessor:
        """
        Get accessor for a specific feature's discovered configs.
        
        Example:
            context.feature(SQSFeature).resources(SQSQueueConfig)
        """
        feature_name = feature_class.__name__
        feature_configs = self.feature_configs.get(feature_name, {})
        
        # Convert string keys back to type objects
        configs_by_type = {}
        for config_type_name, configs in feature_configs.items():
            # Find the actual config class
            for config in configs:
                config_type = type(config)
                if config_type not in configs_by_type:
                    configs_by_type[config_type] = []
                configs_by_type[config_type].extend(configs)
                break
        
        return FeatureAccessor(configs_by_type)

