"""
Step type detection for alignment validation.

Provides functions to detect SageMaker step types and ML frameworks
from scripts and their analysis results.
"""

from typing import List, Optional
from .script_analysis_models import ImportStatement


def detect_step_type_from_registry(script_name: str) -> str:
    """
    Use existing step registry to determine SageMaker step type.
    
    Args:
        script_name: Name of the script (without .py extension)
        
    Returns:
        SageMaker step type (Processing, Training, CreateModel, etc.)
        Defaults to "Processing" if detection fails
    """
    try:
        from cursus.steps.registry.step_names import get_sagemaker_step_type, get_canonical_name_from_file_name
        canonical_name = get_canonical_name_from_file_name(script_name)
        return get_sagemaker_step_type(canonical_name)
    except (ValueError, ImportError, AttributeError):
        # Default fallback to Processing for backward compatibility
        return "Processing"


def detect_framework_from_imports(imports: List) -> Optional[str]:
    """
    Detect framework from existing import analysis.
    
    Args:
        imports: List of ImportStatement objects or strings from script analysis
        
    Returns:
        Detected framework name or None if no framework detected
    """
    framework_patterns = {
        'xgboost': ['xgboost', 'xgb'],
        'pytorch': ['torch', 'pytorch'],
        'sklearn': ['sklearn', 'scikit-learn', 'scikit_learn'],
        'sagemaker': ['sagemaker'],
        'pandas': ['pandas', 'pd'],
        'numpy': ['numpy', 'np']
    }
    
    detected_frameworks = []
    
    for imp in imports:
        # Handle both ImportStatement objects and strings
        if hasattr(imp, 'module_name'):
            module_name_lower = imp.module_name.lower()
        else:
            # Assume it's a string
            module_name_lower = str(imp).lower()
            
        for framework, patterns in framework_patterns.items():
            if any(pattern in module_name_lower for pattern in patterns):
                detected_frameworks.append(framework)
    
    # Return primary framework (prioritize ML frameworks)
    priority_order = ['xgboost', 'pytorch', 'sklearn', 'sagemaker']
    for framework in priority_order:
        if framework in detected_frameworks:
            return framework
    
    # Return first detected framework if no priority match
    return detected_frameworks[0] if detected_frameworks else None


def detect_step_type_from_script_patterns(script_content: str) -> Optional[str]:
    """
    Detect step type from script content patterns.
    
    Args:
        script_content: Content of the script
        
    Returns:
        Detected step type or None if not determinable
    """
    # Training patterns
    training_indicators = [
        'xgb.train(',
        'model.fit(',
        'torch.save(',
        '/opt/ml/model',
        'hyperparameters.json',
        'model.save_model('
    ]
    
    # Processing patterns
    processing_indicators = [
        '/opt/ml/processing/input',
        '/opt/ml/processing/output',
        'pd.read_csv(',
        '.transform(',
        '.fit_transform('
    ]
    
    # CreateModel patterns
    create_model_indicators = [
        'def model_fn(',
        'def input_fn(',
        'def predict_fn(',
        'def output_fn(',
        'pickle.load(',
        'joblib.load('
    ]
    
    # Count pattern matches
    training_score = sum(1 for pattern in training_indicators if pattern in script_content)
    processing_score = sum(1 for pattern in processing_indicators if pattern in script_content)
    create_model_score = sum(1 for pattern in create_model_indicators if pattern in script_content)
    
    # Return type with highest score
    scores = {
        'Training': training_score,
        'Processing': processing_score,
        'CreateModel': create_model_score
    }
    
    max_score = max(scores.values())
    if max_score > 0:
        return max(scores, key=scores.get)
    
    return None


def get_step_type_context(script_name: str, script_content: Optional[str] = None) -> dict:
    """
    Get comprehensive step type context for a script.
    
    Args:
        script_name: Name of the script
        script_content: Optional script content for pattern analysis
        
    Returns:
        Dictionary with step type context information
    """
    context = {
        'script_name': script_name,
        'registry_step_type': None,
        'pattern_step_type': None,
        'final_step_type': None,
        'confidence': 'low'
    }
    
    # Try registry-based detection
    try:
        registry_type = detect_step_type_from_registry(script_name)
        context['registry_step_type'] = registry_type
    except Exception:
        pass
    
    # Try pattern-based detection if script content available
    if script_content:
        try:
            pattern_type = detect_step_type_from_script_patterns(script_content)
            context['pattern_step_type'] = pattern_type
        except Exception:
            pass
    
    # Determine final step type and confidence
    registry_type = context['registry_step_type']
    pattern_type = context['pattern_step_type']
    
    if registry_type and pattern_type:
        if registry_type == pattern_type:
            context['final_step_type'] = registry_type
            context['confidence'] = 'high'
        else:
            # Registry takes precedence but lower confidence
            context['final_step_type'] = registry_type
            context['confidence'] = 'medium'
    elif registry_type:
        context['final_step_type'] = registry_type
        context['confidence'] = 'medium'
    elif pattern_type:
        context['final_step_type'] = pattern_type
        context['confidence'] = 'low'
    else:
        context['final_step_type'] = 'Processing'  # Default
        context['confidence'] = 'low'
    
    return context
