from typing import Union, Tuple, Dict, List, Any

DEFAULT_CONFIG = {
    'device': 'cpu',
    'conf': 0.25,
    'iou': 0.45,
    'classes': None,
    'class_names': None,
    'agnostic': False,
    'fp16': False,
    'replace': None,  # Will be {} if None
    'nc': 0,
    'kpt_shape': (17, 3),
    'pipeline': 'preprocess | inference | postprocess',
    'streaming_mode': False,
    'orig_imgs_shapes': None,  # Will be [] if None
}

def has_nms_in_hef(hef_path):
    """
    check if the .hef file contains postprocess information
    Args:
        hef_path: the path of the .hef file
    
    Returns:
        bool: if the .hef file contains postprocess information
    """ 
    try:
        # Run command and capture output
        import subprocess
        result = subprocess.run(
            ['hailortcli', 'parse-hef', hef_path],
            capture_output=True,
            text=True,
            encoding='utf-8'
        )
        
        # check if "postprocess" is in stdout or stderr
        # use .lower() to prevent case difference (e.g. PostProcess)
        return "postprocess" in result.stdout.lower()
    
    except Exception:
        return False


def parse_pipeline(pipeline: str, stage: str = "preprocess"):
    """
    Parse the pipeline string for the given stage.
    Args:
        pipeline(str): The pipeline string to parse.
        stage(str): The stage to parse the pipeline for.
    Returns:
        list[str]
    Example:
        pipeline = "(letterbox | normalize) | inference | (nms | scale)"
        parse_pipeline(pipeline, "preprocess") -> ["letterbox", "normalize"]
    """
    import re
    idx = {"preprocess": 0, "inference": 1, "postprocess": 2}[stage]
    
    # use regex to split: only match | outside the parentheses
    # logic: find a |, the parentheses after it must be paired
    parts = re.split(r'\s*\|\s*(?![^()]*\))', pipeline)
    
    if idx >= len(parts) or parts[idx] == "None":
        return "None"
    
    # remove parentheses and split content
    content = parts[idx].strip(" ()")
    return [op.strip() for op in content.split('|') if op.strip()]


def load_config(config: Union[str, dict]) -> dict:
    """
    Load configuration from YAML file or dictionary.
    
    This function loads configuration settings from either a YAML file path
    or a dictionary. It also performs special handling for certain configuration
    fields that need type conversion (e.g., kpt_shape, None values).
    
    Args:
        config (str or dict): 
            - str: Path to YAML configuration file (must end with .yaml or .yml)
            - dict: Configuration dictionary (will be copied)
            
    Returns:
        dict: Configuration dictionary with processed values
        
    Raises:
        ValueError: If config type is unsupported or file loading fails
        
    Special Handling:
        - kpt_shape: Converts list to tuple, or extracts numbers from string
        - 'None' strings: Converts string 'None' to Python None value
    """
    import re, yaml
    # Load from YAML file if string path is provided
    if isinstance(config, str) and config.endswith(('yaml', 'yml')):
        try:
            with open(config, 'r') as f:
                setting = yaml.safe_load(f)
        except Exception as e:
            raise ValueError(f"Error loading config: {e}")
    # Use dictionary directly if already a dict
    elif isinstance(config, dict):
        setting = config.copy()
    else:
        raise ValueError(f"Unsupported config type: {type(config)}")
    
    # Special handling: convert kpt_shape from list to tuple
    # kpt_shape (keypoint shape) is used in xpose estimation tasks
    if 'kpt_shape' in setting:
        if isinstance(setting['kpt_shape'], list):
            # Convert list to tuple (required format)
            setting['kpt_shape'] = tuple(setting['kpt_shape'])
        elif isinstance(setting['kpt_shape'], str):
            # Extract numbers from string (e.g., "(17, 3)" -> (17, 3))
            setting['kpt_shape'] = tuple(map(int, re.findall(r'\d+', setting['kpt_shape'])))

    # Special handling: convert string 'None' to Python None
    # YAML files may contain 'None' as a string, which should be converted to None
    for key, value in setting.items():
        if setting[key] == 'None':
            setting[key] = None
        
    return setting


def init_yolo_config(config: Any, streaming_hw: Union[Tuple[int, int], Dict[str, int]] = None) -> dict:
    """
    Initialize and validate YOLO configuration.
    
    Args:
        config: Configuration object, dict, or YAML file path
        streaming_hw: Hardware dimensions for streaming mode
        
    Returns:
        dict: Fully initialized configuration dictionary
    """
    # Step 1: Load config base
    if hasattr(config, '__dict__'):  # Supports YOLOOptions and other objects
        config_dict = {k: v for k, v in config.__dict__.items() if not k.startswith("__")}
    else:
        config_dict = load_config(config)

    # Step 2: Merge default values
    for k, v in DEFAULT_CONFIG.items():
        config_dict.setdefault(k, v)
    
    # Step 2.1: Handle mutable defaults
    if config_dict.get('replace') is None:
        config_dict['replace'] = {}
    if config_dict.get('orig_imgs_shapes') is None:
        config_dict['orig_imgs_shapes'] = []

    # Step 3: Validate required fields
    required_keys = ['version', 'weights', 'task', 'size']
    missing = [k for k in required_keys if config_dict.get(k) in [None, '']]
    if missing:
        raise ValueError(f"Missing required config keys: {missing}")

    # Step 4: Handle streaming mode
    if isinstance(streaming_hw, dict):
        config_dict['orig_imgs_shapes'] = [(streaming_hw['orig_h'], streaming_hw['orig_w'])] * 64
        config_dict['streaming_mode'] = True
    elif isinstance(streaming_hw, (list, tuple)):
        config_dict['orig_imgs_shapes'] = [(streaming_hw[0], streaming_hw[1])] * 64
        config_dict['streaming_mode'] = True
    else:
        config_dict['streaming_mode'] = False
        if not config_dict.get('orig_imgs_shapes'):
            config_dict['orig_imgs_shapes'] = []

    # Step 5: Normalize input shape
    input_shape = config_dict['size']
    if isinstance(input_shape, int):
        config_dict['size'] = (input_shape, input_shape)
    elif isinstance(input_shape, (list, tuple)):
        config_dict['size'] = tuple(input_shape)
    else:
        raise TypeError(f"Unsupported type for input_shape: {type(input_shape).__name__}")
    
    return config_dict


def init_model_info(config: dict, model: Any) -> dict:
    """
    Initialize model-specific configuration like class_names, nc, and kpt_shape.
    
    Args:
        config: Configuration dictionary
        model: Model backend instance (AutoBackend)
        
    Returns:
        dict: Updated configuration dictionary
    """
    # Initialize class_names and nc from model if not provided
    if config.get('class_names') is None:
        if hasattr(model, 'names') and model.names is not None:
            # Get from model
            config['class_names'] = model.names
            if config.get('nc') == 0:
                config['nc'] = len(config['class_names'])
        elif config.get('nc', 0) > 0:
            # Auto-generate names if only nc is known
            config['class_names'] = {i: f'class_{i}' for i in range(config['nc'])}
        else:
            raise ValueError(
                "Failed to initialize: both `nc` and `class_names` are None, "
                "and model does not have `names` attribute."
            )

    # Handle pose task specific settings
    if config.get('task') == 'pose':
        if config.get('kpt_shape') is None:
            config['kpt_shape'] = getattr(model, 'kpt_shape', (17, 3))
            
    return config


def get_imgs_shapes(imgs, img_shape = None, streaming_mode = False, dynamic_hw = None):
    """
    Get the shapes of input images.
    
    Args:
        imgs: Input images
        img_shape: Internal cached image shape (modified in-place)
        streaming_mode: Whether in streaming mode
        dynamic_hw: Dynamic image shape
        
    Returns:
        List of image shapes (h, w)
    """
    import numpy as np
    
    if dynamic_hw:
        return [dynamic_hw] * len(imgs)

    if streaming_mode or (isinstance(imgs, np.ndarray) and img_shape):
        return img_shape

    if isinstance(imgs, np.ndarray):
        s = imgs.shape
        h, w = (s[1], s[2]) if s[1] > 3 else (s[2], s[3])
        shapes = [(h, w)] * s[0]
        if isinstance(img_shape, list) and not img_shape:
            img_shape.extend(shapes)
        return shapes

    return [i.shape[:2] for i in imgs]


def timer():
    def decorator(func):
        def wrapper(*args, **kwargs):
            self = args[0]  # Ensure one of them is `self`
            if hasattr(self, 'verbose') and self.verbose:
                import time
                import torch
                
                is_cuda = (hasattr(self, 'device') and
                           isinstance(self.device, str) and
                           self.device.split(':')[0] == 'cuda')
                if is_cuda:
                    torch.cuda.synchronize()
                start_time = time.time()
                result = func(*args, **kwargs)
                if is_cuda:
                    torch.cuda.synchronize()
                end_time = time.time()
                print(f"{func.__name__}: {(end_time - start_time)*1000:.1f} ms")
            else:
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator
