"""
後處理 Pipeline Builder
"""
from typing import List, Tuple, Union, Callable
import numpy as np

from .utils.common import has_nms_in_hef, parse_pipeline
from .utils.postprocess import (
    # Constants
    TASK_ALLOWED_STEPS,
    TASK_DEFAULT_STEPS,
    # Core functions
    nms,
    unpack,
    scale,
    mask,
    keypoints,
    engine_swap,
    proto,
    classify,
)

# IO signature for each step: (required_input_fields, output_fields)
# Both fields are frozensets of named state keys tracked during pipeline validation.
# nms has task-specific IO, represented as a nested dict keyed by task name.
#
# To add a new step:
#   1. Add an entry here declaring what fields it needs and what it produces.
#   2. Add the corresponding wrapper lambda in _create_step_wrappers().
#   3. Add the step name to TASK_ALLOWED_STEPS (and TASK_DEFAULT_STEPS if applicable)
#      in pipeline/utils/postprocess.py.
_POSTPROCESS_STEP_IO = {
    'unpack':      (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
    'engine_swap': (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
    'proto':       (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes', 'proto_output'])),
    'nms': {
        'detect':  (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
        'segment': (frozenset(['preds', 'orig_imgs_shapes', 'proto_output']),
                   frozenset(['preds', 'orig_imgs_shapes', 'proto_output'])),
        'pose':    (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
        'classify':(frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
    },
    'mask':        (frozenset(['preds', 'orig_imgs_shapes', 'proto_output']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
    'scale':       (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
    'keypoints':   (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
    'classify':    (frozenset(['preds', 'orig_imgs_shapes']),
                   frozenset(['preds', 'orig_imgs_shapes'])),
}


def _validate_postprocess_pipeline(steps: list, task: str) -> None:
    """Statically validate that pipeline steps can chain together at construction time, raising an error early if not."""
    current_fields = frozenset(['preds', 'orig_imgs_shapes'])
    for step in steps:
        sig = _POSTPROCESS_STEP_IO.get(step)
        if sig is None:
            raise ValueError(f"Unknown postprocess step: '{step}'")
        if isinstance(sig, dict):
            if task not in sig:
                raise ValueError(f"Step '{step}' does not support task '{task}'")
            sig = sig[task]
        required, output = sig
        missing = required - current_fields
        if missing:
            raise ValueError(
                f"Pipeline steps cannot be chained: '{step}' requires fields {set(missing)}, "
                f"but the current state only contains {set(current_fields)}.\n"
                f"  → Ensure a step that produces these fields (e.g. 'proto') appears before '{step}'."
            )
        current_fields = output



class PostProcessPipeline:
    """
    Postprocessing pipeline for YOLO models.
    Supports dynamic updating of parameters like confidence and IoU thresholds.

    Usage Examples:
        # 1. Dynamic update via YOLO instance
        model.set_postprocess_params(conf=0.6, iou=0.5)
        
        # 2. Temporary pipeline structure override
        model.set_postprocess_params(pipeline="scale")
        
        # 3. Restore to original config settings
        model.set_postprocess_params(conf="original", iou="original", pipeline="original")
    """
    def __init__(self, config: dict):
        self.config = config
        
        # Initialize parameters from config
        self.conf = config.get('conf', 0.25)
        self.iou = config.get('iou', 0.45)
        self.nc = config.get('nc', 80)
        self.classes = config.get('classes')
        self.agnostic = config.get('agnostic', False)
        self.task = config['task']
        self.weights = config['weights']
        self.version = config.get('version', 'yolo11')
        self.kpt_shape = config.get('kpt_shape', (17, 3))
        
        # 1. Normalize input shape
        input_shape = config['size']
        if isinstance(input_shape, int):
            self.input_shape = (input_shape, input_shape)
        else:
            self.input_shape = tuple(input_shape)
        self.original_input_shape = self.input_shape
            
        # 2. Store original config values for restoration
        self.original_conf = self.conf
        self.original_iou = self.iou
        self.original_classes = self.classes
        self.original_agnostic = self.agnostic

        # 3. Parse and build pipeline steps
        self.original_pipeline = config.get('pipeline', '')
        self.internal_steps = self._determine_steps(self.original_pipeline)
        self.step_wrappers = self._create_step_wrappers()

    def _determine_steps(self, pipeline_str: str) -> List[str]:
        """Determine which processing steps to execute based on pipeline string."""
        # Auto-detect built-in NMS for .hef files
        if self.weights.endswith('.hef'):
            hef_has_nms = has_nms_in_hef(self.weights)
        else:
            hef_has_nms = False
            
        user_steps = parse_pipeline(pipeline_str, stage="postprocess")
        
        if user_steps == "None" or user_steps is None:
            return []
            
        if isinstance(user_steps, list) and len(user_steps) > 0:
            user_steps_lower = [s.strip().lower() for s in user_steps]
            if 'postprocess' in user_steps_lower:
                steps = TASK_DEFAULT_STEPS.get(self.task, []).copy()
                if self.task == 'detect' and self.weights.endswith('.hef'):
                    steps = ['unpack', 'scale'] if hef_has_nms else ['nms', 'scale']
            else:
                steps = user_steps_lower
        else:
            # Default steps
            steps = TASK_DEFAULT_STEPS.get(self.task, []).copy()
            if self.task == 'detect' and self.weights.endswith('.hef'):
                steps = ['unpack', 'scale'] if hef_has_nms else ['nms', 'scale']

        # Validate step names
        allowed = TASK_ALLOWED_STEPS.get(self.task, [])
        for step in steps:
            if step not in allowed:
                raise ValueError(f"Step '{step}' is not allowed for task '{self.task}'. Allowed: {allowed}")

        # Validate pipeline IO compatibility
        _validate_postprocess_pipeline(steps, self.task)

        return steps

    def set_pipeline(self, pipeline_str: Union[str, None] = None):
        """Dynamically switch the pipeline steps. Pass 'original' to reset."""
        if pipeline_str == 'original' or pipeline_str is None:
            pipeline_str = self.original_pipeline
            
        self.internal_steps = self._determine_steps(pipeline_str)
        self.step_wrappers = self._create_step_wrappers()

    def update(self, **kwargs):
        """Unified update method for speed and special logic."""
        if 'pipeline' in kwargs:
            self.set_pipeline(kwargs.pop('pipeline'))
            
        if 'conf' in kwargs and kwargs['conf'] == 'original':
            kwargs['conf'] = self.original_conf
            
        if 'iou' in kwargs and kwargs['iou'] == 'original':
            kwargs['iou'] = self.original_iou

        if 'classes' in kwargs and kwargs['classes'] == 'original':
            kwargs['classes'] = self.original_classes

        if 'agnostic' in kwargs and kwargs['agnostic'] == 'original':
            kwargs['agnostic'] = self.original_agnostic

        if 'size' in kwargs:
            size = kwargs.pop('size')
            if size == 'original':
                self.input_shape = self.original_input_shape
            else:
                self.input_shape = (size, size) if isinstance(size, int) else tuple(size)
            
        # Bulk update remaining parameters for speed
        _ALLOWED_KEYS = {'conf', 'iou', 'classes', 'agnostic', 'nc', 'version', 'kpt_shape'}
        invalid = set(kwargs) - _ALLOWED_KEYS
        if invalid:
            raise ValueError(f"Unknown parameter(s) for PostProcessPipeline.update(): {invalid}")
        self.__dict__.update(kwargs)

    def _create_step_wrappers(self) -> List[Callable]:
        """Create function wrappers for each internal step, referencing self for dynamic params."""
        wrappers = []
        
        for step_name in self.internal_steps:
            if step_name == 'unpack':
                wrappers.append(
                    lambda preds, orig_imgs_shapes: unpack(preds, self.input_shape)
                )
            
            elif step_name == 'nms':
                if self.task == 'detect':
                    wrappers.append(
                        lambda preds, orig_imgs_shapes: 
                            nms(preds[0], self.conf, self.iou, self.nc, self.classes, self.agnostic, self.version)
                    )
                elif self.task == 'segment':
                    wrappers.append(
                        lambda state: (
                            nms(state[0][0], self.conf, self.iou, self.nc, self.classes, self.agnostic, self.version, self.task),
                            state[1],
                            state[2]
                        )
                    )
                elif self.task == 'pose':
                    if self.weights.endswith('.engine'):
                        wrappers.append(
                            lambda preds, orig_imgs_shapes: [preds[-1]]
                        )
                    wrappers.append(
                        lambda preds, orig_imgs_shapes: 
                            nms(preds[0], self.conf, self.iou, self.nc, self.classes, self.agnostic, self.version, self.task)
                    )
            
            elif step_name == 'scale':
                wrappers.append(
                    lambda preds, orig_imgs_shapes: scale(preds, self.input_shape, orig_imgs_shapes)
                )
            
            elif step_name == 'engine_swap':
                wrappers.append(
                    lambda preds, orig_imgs_shapes: (engine_swap(preds, self.weights), orig_imgs_shapes)
                )
            
            elif step_name == 'proto':
                wrappers.append(
                    lambda preds, orig_imgs_shapes: (preds, orig_imgs_shapes, proto(preds))
                )
            
            elif step_name == 'mask':
                wrappers.append(
                    lambda state: 
                        mask(state[0], state[2], self.input_shape, state[1], self.version)
                )
            
            elif step_name == 'keypoints':
                wrappers.append(
                    lambda preds, orig_imgs_shapes: 
                        keypoints(preds, self.input_shape, orig_imgs_shapes, self.kpt_shape)
                )
            
            elif step_name == 'classify':
                if self.weights.endswith('.engine'):
                    wrappers.append(
                        lambda preds, orig_imgs_shapes: [preds[-1]]
                    )
                wrappers.append(
                    lambda preds, orig_imgs_shapes: classify(preds)
                )
        
        return wrappers

    def __call__(self, preds: Union[List[np.ndarray], Tuple], orig_imgs_shapes: List[Tuple[int, int]]):
        """Execute the postprocessing pipeline."""
        if not self.step_wrappers:
            return preds

        # State starts as a 2-tuple (preds, orig_imgs_shapes).
        # The dispatch rule is:
        #   - tuple state  → unpack with *state (wrapper receives individual positional args)
        #   - non-tuple    → pass as first arg alongside the original orig_imgs_shapes
        #
        # Wrappers that need to carry extra data (e.g. proto_output for segment) return
        # a tuple to keep that data flowing; wrappers that produce a final list (e.g. nms
        # for detect) return a non-tuple, switching to the else branch for subsequent steps.
        state = (preds, orig_imgs_shapes)
        for wrapper in self.step_wrappers:
            if isinstance(state, tuple):
                state = wrapper(*state)
            else:
                state = wrapper(state, orig_imgs_shapes)

        return state


def build_postprocess_pipeline(config: dict) -> PostProcessPipeline:
    """Factory function to build the postprocessing pipeline class."""
    return PostProcessPipeline(config)

