import logging
from pathlib import Path
from typing import Dict, Optional, List, Tuple

import numpy as np
import cv2 # For potential direct cv2 operations if ipu doesn't cover all merge needs

from .base_stage import ProcessingStage
from ..asset_context import AssetProcessingContext
from rule_structure import FileRule
from utils.path_utils import sanitize_filename
from ...utils import image_processing_utils as ipu


logger = logging.getLogger(__name__)

class MapMergingStage(ProcessingStage):
    """
    Merges individually processed maps based on MAP_MERGE rules.
    This stage performs operations like channel packing.
    """

    def execute(self, context: AssetProcessingContext) -> AssetProcessingContext:
        """
        Executes the map merging logic.

        Args:
            context: The asset processing context.

        Returns:
            The updated asset processing context.
        """
        asset_name_for_log = context.asset_rule.asset_name if context.asset_rule else "Unknown Asset"
        if context.status_flags.get('skip_asset'):
            logger.info(f"Skipping map merging for asset {asset_name_for_log} as skip_asset flag is set.")
            return context

        if not hasattr(context, 'merged_maps_details'):
            context.merged_maps_details = {}
        
        if not hasattr(context, 'processed_maps_details'):
            logger.warning(f"Asset {asset_name_for_log}: 'processed_maps_details' not found in context. Cannot perform map merging.")
            return context

        if not context.files_to_process: # This list might not be relevant if merge rules are defined elsewhere or implicitly
            logger.info(f"Asset {asset_name_for_log}: No files_to_process defined. This stage might rely on config or processed_maps_details directly for merge rules.")
            # Depending on design, this might not be an error, so we don't return yet.

        logger.info(f"Starting MapMergingStage for asset: {asset_name_for_log}")

        # TODO: The logic for identifying merge rules and their inputs needs significant rework
        # as FileRule no longer has 'id' or 'merge_settings' directly in the way this stage expects.
        # Merge rules are likely defined in the main configuration (context.config_obj.map_merge_rules)
        # and need to be matched against available maps in context.processed_maps_details.

        # Placeholder for the loop that would iterate over context.config_obj.map_merge_rules
        # For now, this stage will effectively do nothing until that logic is implemented.
        
        # Example of how one might start to adapt:
        # for configured_merge_rule in context.config_obj.map_merge_rules:
        #     output_map_type = configured_merge_rule.get('output_map_type')
        #     inputs_config = configured_merge_rule.get('inputs') # e.g. {"R": "NORMAL", "G": "ROUGHNESS"}
        #     # ... then find these input map_types in context.processed_maps_details ...
        #     # ... and perform the merge ...
        #     # This is a complex change beyond simple attribute renaming.

        # The following is the original loop structure, which will likely fail due to missing attributes on FileRule.
        # Keeping it commented out to show what was there.
        """
        for merge_rule in context.files_to_process: # This iteration logic is likely incorrect for merge rules
            if not isinstance(merge_rule, FileRule) or merge_rule.item_type != "MAP_MERGE":
                continue

            # FileRule does not have merge_settings or id.hex
            # This entire block needs to be re-thought based on where merge rules are defined.
            # Assuming merge_rule_id_hex would be a generated UUID for this operation.
            merge_rule_id_hex = f"merge_op_{uuid.uuid4().hex[:8]}"
            current_map_type = merge_rule.item_type_override or merge_rule.item_type

            logger.error(f"Asset {asset_name_for_log}, Potential Merge for {current_map_type}: Merge rule processing needs rework. FileRule lacks 'merge_settings' and 'id'. Skipping this rule.")
            context.merged_maps_details[merge_rule_id_hex] = {
                'map_type': current_map_type,
                'status': 'Failed',
                'reason': 'Merge rule processing logic in MapMergingStage needs refactor due to FileRule changes.'
            }
            continue
        """

        # For now, let's assume no merge rules are processed until the logic is fixed.
        num_merge_rules_attempted = 0
        # If context.config_obj.map_merge_rules exists, iterate it here.
        # The original code iterated context.files_to_process looking for item_type "MAP_MERGE".
        # This implies FileRule objects were being used to define merge operations, which is no longer the case
        # if 'merge_settings' and 'id' were removed from FileRule.

        # The core merge rules are in context.config_obj.map_merge_rules
        # Each rule in there defines an output_map_type and its inputs.
        
        config_merge_rules = context.config_obj.map_merge_rules
        if not config_merge_rules:
            logger.info(f"Asset {asset_name_for_log}: No map_merge_rules found in configuration. Skipping map merging.")
            return context

        for rule_idx, configured_merge_rule in enumerate(config_merge_rules):
            output_map_type = configured_merge_rule.get('output_map_type')
            inputs_map_type_to_channel = configured_merge_rule.get('inputs') # e.g. {"R": "NRM", "G": "NRM", "B": "ROUGH"}
            default_values = configured_merge_rule.get('defaults', {}) # e.g. {"R": 0.5, "G": 0.5, "B": 0.5}
            # output_bit_depth_rule = configured_merge_rule.get('output_bit_depth', 'respect_inputs') # Not used yet

            if not output_map_type or not inputs_map_type_to_channel:
                logger.warning(f"Asset {asset_name_for_log}: Invalid configured_merge_rule at index {rule_idx}. Missing 'output_map_type' or 'inputs'. Rule: {configured_merge_rule}")
                continue
            
            num_merge_rules_attempted +=1
            merge_op_id = f"merge_{sanitize_filename(output_map_type)}_{rule_idx}"
            logger.info(f"Asset {asset_name_for_log}: Processing configured merge rule for '{output_map_type}' (Op ID: {merge_op_id})")

            loaded_input_maps: Dict[str, np.ndarray] = {} # Key: input_map_type (e.g. "NRM"), Value: image_data
            input_map_paths: Dict[str, str] = {} # Key: input_map_type, Value: path_str
            target_dims: Optional[Tuple[int, int]] = None
            all_inputs_valid = True
            
            # Find and load input maps from processed_maps_details
            # This assumes one processed map per map_type. If multiple variants exist, this needs refinement.
            required_input_map_types = set(inputs_map_type_to_channel.values())
            
            for required_map_type in required_input_map_types:
                found_processed_map = None
                processed_map_key = None
                for p_key, p_details in context.processed_maps_details.items():
                    processed_map_type_in_details = p_details.get('map_type')
                    # Check for direct match or match with "MAP_" prefix
                    if (processed_map_type_in_details == required_map_type or \
                        processed_map_type_in_details == f"MAP_{required_map_type}") and \
                       p_details.get('status') == 'Processed':
                        found_processed_map = p_details
                        processed_map_key = p_key # The UUID hex key from individual processing
                        break
                
                if not found_processed_map:
                    logger.warning(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Required input map_type '{required_map_type}' for output '{output_map_type}' not found or not processed in processed_maps_details.")
                    # Option: Use default value for the entire map if one could be constructed for this map_type
                    # For now, we fail the merge if a required map is missing.
                    all_inputs_valid = False
                    context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': f"Required input map_type '{required_map_type}' missing."}
                    break # Break from finding inputs for this merge rule

                temp_file_path = Path(found_processed_map['temp_processed_file'])
                if not temp_file_path.exists():
                    logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Temp file {temp_file_path} for input map_type '{required_map_type}' does not exist.")
                    all_inputs_valid = False
                    context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': f"Temp file for input '{required_map_type}' missing."}
                    break

                try:
                    image_data = ipu.load_image(temp_file_path)
                    if image_data is None: raise ValueError("Loaded image is None")
                except Exception as e:
                    logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Error loading image {temp_file_path} for input map_type '{required_map_type}': {e}")
                    all_inputs_valid = False
                    context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': f"Error loading input '{required_map_type}'."}
                    break
                
                loaded_input_maps[required_map_type] = image_data
                input_map_paths[required_map_type] = str(temp_file_path)

                current_dims = (image_data.shape[1], image_data.shape[0])
                if target_dims is None:
                    target_dims = current_dims
                elif current_dims != target_dims:
                    logger.warning(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Input map '{required_map_type}' dims {current_dims} differ from target {target_dims}. Resizing.")
                    try:
                        image_data = ipu.resize_image(image_data, target_dims[0], target_dims[1])
                        if image_data is None: raise ValueError("Resize returned None")
                        loaded_input_maps[required_map_type] = image_data
                    except Exception as e:
                        logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Failed to resize '{required_map_type}': {e}")
                        all_inputs_valid = False
                        context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': f"Failed to resize input '{required_map_type}'."}
                        break
            
            if not all_inputs_valid:
                logger.warning(f"Asset {asset_name_for_log}: Skipping merge for Op ID {merge_op_id} ('{output_map_type}') due to invalid inputs.")
                continue

            if not loaded_input_maps or target_dims is None:
                logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: No input maps loaded or target_dims not set for '{output_map_type}'. This shouldn't happen if all_inputs_valid was true.")
                context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': 'Internal error: input maps not loaded or target_dims missing.'}
                continue

            # Determine output channels (e.g., 3 for RGB, 1 for Grayscale)
            # This depends on the keys in inputs_map_type_to_channel (R,G,B,A)
            output_channel_keys = sorted(list(inputs_map_type_to_channel.keys())) # e.g. ['B', 'G', 'R']
            num_output_channels = len(output_channel_keys)
            
            if num_output_channels == 0:
                logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: No output channels defined in 'inputs' for '{output_map_type}'.")
                context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': 'No output channels defined.'}
                continue

            try:
                if num_output_channels == 1: # Grayscale output
                    merged_image = np.zeros((target_dims[1], target_dims[0]), dtype=np.uint8)
                else: # Color output
                    merged_image = np.zeros((target_dims[1], target_dims[0], num_output_channels), dtype=np.uint8)
            except Exception as e:
                logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Error creating empty merged image for '{output_map_type}': {e}")
                context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': f'Error creating output canvas: {e}'}
                continue

            merge_op_failed_detail = False
            for i, out_channel_char in enumerate(output_channel_keys): # e.g. R, G, B
                input_map_type_for_this_channel = inputs_map_type_to_channel[out_channel_char]
                source_image = loaded_input_maps.get(input_map_type_for_this_channel)
                
                source_data_this_channel = None
                if source_image is not None:
                    if source_image.ndim == 2: # Grayscale source
                        source_data_this_channel = source_image
                    elif source_image.ndim == 3 or source_image.ndim == 4: # Color source (3-channel BGR or 4-channel BGRA), assumed loaded by ipu.load_image
                        # Standard BGR(A) channel indexing: B=0, G=1, R=2, A=3 (if present)
                        # This map helps get NRM's Red data for 'R' output, NRM's Green for 'G' output etc.
                        # based on the semantic meaning of out_channel_char.
                        semantic_to_bgr_idx = {'R': 2, 'G': 1, 'B': 0, 'A': 3}

                        if input_map_type_for_this_channel == "NRM":
                            idx_to_extract = semantic_to_bgr_idx.get(out_channel_char)
                            
                            if idx_to_extract is not None and idx_to_extract < source_image.shape[2]:
                                source_data_this_channel = source_image[:, :, idx_to_extract]
                                logger.debug(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: For output '{out_channel_char}', using NRM's semantic '{out_channel_char}' channel (BGR(A) index {idx_to_extract}).")
                            else:
                                # Fallback if out_channel_char isn't R,G,B,A or NRM doesn't have the channel (e.g. 3-channel NRM and 'A' requested)
                                logger.warning(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Could not map output '{out_channel_char}' to a specific BGR(A) channel of NRM (shape {source_image.shape}). Defaulting to NRM's channel 0 (Blue).")
                                source_data_this_channel = source_image[:, :, 0]
                        else:
                            # For other multi-channel sources (e.g., ROUGH as RGB, or other color maps not "NRM")
                            # Default to taking the first channel (Blue in BGR).
                            # This covers "Roughness map's greyscale data" if ROUGH is RGB (by taking one of its channels as a proxy).
                            source_data_this_channel = source_image[:, :, 0]
                            logger.debug(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: For output '{out_channel_char}', source {input_map_type_for_this_channel} (shape {source_image.shape}) is multi-channel but not NRM. Using its channel 0 (Blue).")
                else: # Source map was not found, use default
                    default_val_for_channel = default_values.get(out_channel_char)
                    if default_val_for_channel is not None:
                        # Convert 0-1 float default to 0-255 uint8
                        source_data_this_channel = np.full((target_dims[1], target_dims[0]), int(default_val_for_channel * 255), dtype=np.uint8)
                        logger.info(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Using default value {default_val_for_channel} for output channel '{out_channel_char}' as input map '{input_map_type_for_this_channel}' was missing.")
                    else:
                        logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Input map '{input_map_type_for_this_channel}' for output channel '{out_channel_char}' is missing and no default value provided.")
                        merge_op_failed_detail = True; break
                
                if source_data_this_channel is None: # Should be caught by default value logic or earlier checks
                    logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Failed to get source data for output channel '{out_channel_char}'.")
                    merge_op_failed_detail = True; break

                try:
                    if merged_image.ndim == 2: # Single channel output
                        merged_image = source_data_this_channel
                    else: # Multi-channel output
                        merged_image[:, :, i] = source_data_this_channel
                except Exception as e:
                    logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Error assigning data to output channel '{out_channel_char}' (index {i}): {e}")
                    merge_op_failed_detail = True; break
            
            if merge_op_failed_detail:
                context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': 'Error during channel assignment.'}
                continue

            output_format = 'png' # Default, can be configured per rule later
            temp_merged_filename = f"merged_{sanitize_filename(output_map_type)}_{merge_op_id}.{output_format}"
            temp_merged_path = context.engine_temp_dir / temp_merged_filename

            try:
                save_success = ipu.save_image(temp_merged_path, merged_image)
                if not save_success: raise ValueError("Save image returned false")
            except Exception as e:
                logger.error(f"Asset {asset_name_for_log}, Merge Op ID {merge_op_id}: Error saving merged image {temp_merged_path}: {e}")
                context.merged_maps_details[merge_op_id] = {'map_type': output_map_type, 'status': 'Failed', 'reason': f'Failed to save merged image: {e}'}
                continue
            
            logger.info(f"Asset {asset_name_for_log}: Successfully merged and saved '{output_map_type}' (Op ID: {merge_op_id}) to {temp_merged_path}")
            context.merged_maps_details[merge_op_id] = {
                'map_type': output_map_type,
                'temp_merged_file': str(temp_merged_path),
                'input_map_types_used': list(inputs_map_type_to_channel.values()),
                'input_map_files_used': input_map_paths,
                'merged_dimensions': target_dims,
                'status': 'Processed'
            }

        logger.info(f"Finished MapMergingStage for asset: {asset_name_for_log}. Merged maps operations attempted: {num_merge_rules_attempted}, Succeeded: {len([d for d in context.merged_maps_details.values() if d.get('status') == 'Processed'])}")
        return context