import numpy as np

from .utils.common import parse_pipeline
from .utils.preprocess import (
    PREPROCESS_ALLOWED_STEPS,
    PREPROCESS_DEFAULT_STEPS,
    letterbox_batch,
    list2arr,
    bgr2rgb,
    hwc2chw,
    normalize,
    contiguous,
    letterbox_classify_batch,
)
from .utils.ops import letterbox_classify

# IO signature for each step: (accepted_input_formats, output_format)
# Format tags describe the array layout at each stage:
#   'list'  — Python list of images (before any stacking)
#   'bhwc'  — ndarray (N, H, W, C), channels-last
#   'bchw'  — ndarray (N, C, H, W), channels-first (after hwc2chw)
#
# To add a new step:
#   1. Add an entry here with the set of accepted input formats and the output format.
#   2. Add the corresponding function/lambda in _init_step_map().
#   3. Add the step name to PREPROCESS_ALLOWED_STEPS (and PREPROCESS_DEFAULT_STEPS
#      if applicable) in pipeline/utils/preprocess.py.
_PREPROCESS_STEP_IO = {
    'list2arr':   ({'list'},           'bhwc'),
    'letterbox':  ({'bhwc'},           'bhwc'),
    'bgr2rgb':    ({'bhwc'},           'bhwc'),
    'hwc2chw':    ({'bhwc'},           'bchw'),
    'normalize':  ({'bchw'},           'bchw'),
    'contiguous': ({'bhwc', 'bchw'},   'bchw'),
    'classify':   ({'list', 'bhwc'},   'bchw'),
}


def _validate_preprocess_pipeline(steps: list) -> None:
    """Statically validate that pipeline steps can chain together at construction time, raising an error early if not."""
    if not steps:
        return
    initial_format = 'list' if steps[0] == 'list2arr' else 'bhwc'
    current_format = initial_format
    for step in steps:
        sig = _PREPROCESS_STEP_IO.get(step)
        if sig is None:
            raise ValueError(f"Unknown preprocess step: '{step}'")
        accepted_inputs, output_format = sig
        if current_format not in accepted_inputs:
            raise ValueError(
                f"Pipeline steps cannot be chained: '{step}' accepts formats {accepted_inputs}, "
                f"but the current format is '{current_format}'.\n"
                f"  → Common cause: steps requiring HWC format (e.g. 'bgr2rgb', 'letterbox') "
                f"cannot follow 'hwc2chw'."
            )
        current_format = output_format


# ===== Pipeline Builder =====
class PreProcessPipeline:
    """
    Preprocessing pipeline for YOLO models.
    Supports dynamic updating of parameters like input size and step sequence.

    Usage Examples:
        # 1. Dynamic update via YOLO instance (Recommended)
        model.set_preprocess_params(pipeline="letterbox | bgr2rgb", size=640)
        model.set_preprocess_params(pipeline="letterbox | bgr2rgb")
        model.set_preprocess_params(size=640)
        
        # 2. Results in temporary override:
        # - Only 'letterbox' and 'bgr2rgb' steps will execute.
        # - Image will be resized to 640x640.
        
        # 3. Restore to original config settings:
        model.set_preprocess_params(pipeline="original", size="original")
    """
    
    def __init__(self, config: dict):
        self.config = config
        self.task = config['task']
        self.weights = config['weights']
        
        # 1. Normalize and store input shape
        self._set_input_shape(config['size'])
        self.original_input_shape = self.input_shape
        
        # 2. Initialize step map (lambdas reference self for dynamic params)
        self._init_step_map()
        
        # 3. Store original steps and initialize active steps
        self.original_steps = self._determine_steps(self.config.get('pipeline', ''))
        self.active_steps = list(self.original_steps)
        self.step_wrappers = [self.step_map[s] for s in self.active_steps]

    def _set_input_shape(self, size):
        """Helper to normalize size into (H, W)."""
        if isinstance(size, int):
            self.input_shape = (size, size)
        else:
            self.input_shape = tuple(size)

    def set_size(self, size):
        """Dynamically update input size. Pass 'original' to reset."""
        if size == 'original':
            self.input_shape = self.original_input_shape
        else:
            self._set_input_shape(size)

    def update(self, **kwargs):
        """Unified update method for speed and special logic."""
        if 'pipeline' in kwargs:
            self.set_pipeline(kwargs.pop('pipeline'))
            
        if 'size' in kwargs:
            self.set_size(kwargs.pop('size'))
            
        # Bulk update remaining parameters for speed
        if kwargs:
            raise ValueError(f"Unknown parameter(s) for PreProcessPipeline.update(): {set(kwargs)}")

    def _init_step_map(self):
        """Map step names to their corresponding functions/wrappers."""
        self.step_map = {
            'list2arr': list2arr,
            'letterbox': lambda imgs: letterbox_batch(imgs, self.input_shape),
            'bgr2rgb': bgr2rgb,
            'hwc2chw': hwc2chw,
            'normalize': normalize,
            'contiguous': contiguous,
            'classify': lambda imgs: letterbox_classify_batch(
                imgs, lambda im: letterbox_classify(im, size=self.input_shape)
            )
        }

    def _determine_steps(self, pipeline_str: str) -> list:
        """Determine which processing steps to execute based on a pipeline string."""
        user_steps = parse_pipeline(pipeline_str, stage="preprocess")
        
        if user_steps == "None":
            return []
            
        if isinstance(user_steps, list) and len(user_steps) > 0:
            user_steps_lower = [s.strip().lower() for s in user_steps]
            if 'preprocess' in user_steps_lower:
                steps = PREPROCESS_DEFAULT_STEPS.get(self.task, []).copy()
            else:
                steps = user_steps_lower
        else:
            steps = PREPROCESS_DEFAULT_STEPS.get(self.task, []).copy()
            
        # Adjust for .hef and .axmodel
        if self.weights.endswith(('.hef', '.axmodel')):
            if 'hwc2chw' in steps:
                steps.remove('hwc2chw')
            if 'normalize' in steps:
                steps.remove('normalize')

        # Validate pipeline IO compatibility
        _validate_preprocess_pipeline(steps)

        return steps

    def set_pipeline(self, pipeline_str = None):
        """Dynamically switch the pipeline steps. Pass 'original' to reset to original config."""
        if pipeline_str == 'original' or pipeline_str is None:
            self.active_steps = list(self.original_steps)
        else:
            self.active_steps = self._determine_steps(pipeline_str)
        self.step_wrappers = [self.step_map[s] for s in self.active_steps]

    def __call__(self, imgs):
        """Execute the preprocessing pipeline."""
        if not self.step_wrappers:
            return imgs
            
        # Allow single image input
        single = False
        if isinstance(imgs, np.ndarray) and imgs.ndim == 3:
            imgs = imgs[None, ...]  # add batch dim
            single = True

        # Apply steps
        for fn in self.step_wrappers:
            imgs = fn(imgs)

        return imgs[0] if single else imgs


def build_preprocess_pipeline(config: dict) -> PreProcessPipeline:
    """Factory function to build the preprocessing pipeline class."""
    return PreProcessPipeline(config)
