import logging
from pathlib import Path
import numpy as np
from typing import List

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

logger = logging.getLogger(__name__)

class GlossToRoughConversionStage(ProcessingStage):
    """
    Processing stage to convert glossiness maps to roughness maps.
    Iterates through FileRules, identifies GLOSS maps, loads their
    corresponding temporary processed images, inverts them, and saves
    them as new temporary ROUGHNESS maps. Updates the FileRule and
    context.processed_maps_details accordingly.
    """

    def execute(self, context: AssetProcessingContext) -> AssetProcessingContext:
        """
        Executes the gloss to roughness conversion logic.

        Args:
            context: The AssetProcessingContext containing asset and processing details.

        Returns:
            The updated AssetProcessingContext.
        """
        asset_name_for_log = context.asset_rule.asset_name if context.asset_rule else "Unknown Asset"
        if context.status_flags.get('skip_asset'):
            logger.debug(f"Asset '{asset_name_for_log}': Skipping GlossToRoughConversionStage due to skip_asset flag.")
            return context

        if not context.files_to_process or not context.processed_maps_details:
            logger.debug(
                f"Asset '{asset_name_for_log}': No files to process or processed_maps_details empty "
                f"in GlossToRoughConversionStage. Skipping."
            )
            return context

        new_files_to_process: List[FileRule] = []
        processed_a_gloss_map = False

        logger.info(f"Asset '{asset_name_for_log}': Starting Gloss to Roughness Conversion Stage.")

        for idx, file_rule in enumerate(context.files_to_process):
            # Assuming FileRule has 'map_type' and 'id' (with a .hex attribute) and 'source_file_path'
            # These might need to be checked with hasattr if they are optional or could be missing
            if hasattr(file_rule, 'map_type') and file_rule.map_type == "GLOSS":
                if not hasattr(file_rule, 'id') or not hasattr(file_rule.id, 'hex'):
                    logger.warning(f"Asset '{asset_name_for_log}': GLOSS FileRule missing 'id.hex'. Skipping conversion for this rule: {file_rule}")
                    new_files_to_process.append(file_rule)
                    continue
                map_detail_key = file_rule.id.hex
                
                source_file_path_for_log = file_rule.source_file_path if hasattr(file_rule, 'source_file_path') else "Unknown source path"

                if map_detail_key not in context.processed_maps_details:
                    logger.warning(
                        f"Asset '{asset_name_for_log}': GLOSS map '{source_file_path_for_log}' "
                        f"(ID: {map_detail_key}) found in files_to_process but not in processed_maps_details. "
                        f"Adding original rule and skipping conversion for this map."
                    )
                    new_files_to_process.append(file_rule)
                    continue

                map_details = context.processed_maps_details[map_detail_key]
                
                if map_details.get('status') != 'Processed' or 'temp_processed_file' not in map_details:
                    logger.warning(
                        f"Asset '{asset_name_for_log}': GLOSS map '{source_file_path_for_log}' "
                        f"(ID: {map_detail_key}) not successfully processed by previous stage or temp file missing. "
                        f"Status: {map_details.get('status')}. Adding original rule and skipping conversion."
                    )
                    new_files_to_process.append(file_rule)
                    continue

                original_temp_path_str = map_details['temp_processed_file']
                original_temp_path = Path(original_temp_path_str)

                if not original_temp_path.exists():
                    logger.error(
                        f"Asset '{asset_name_for_log}': Temporary file {original_temp_path_str} for GLOSS map "
                        f"(ID: {map_detail_key}) does not exist. Adding original rule and skipping conversion."
                    )
                    new_files_to_process.append(file_rule)
                    continue
                
                logger.debug(f"Asset '{asset_name_for_log}': Processing GLOSS map {original_temp_path} for conversion.")
                image_data = ipu.load_image(original_temp_path)

                if image_data is None:
                    logger.error(
                        f"Asset '{asset_name_for_log}': Failed to load image data from {original_temp_path} "
                        f"for GLOSS map (ID: {map_detail_key}). Adding original rule and skipping conversion."
                    )
                    new_files_to_process.append(file_rule)
                    continue

                # Perform Inversion
                inverted_image_data: np.ndarray
                if np.issubdtype(image_data.dtype, np.floating):
                    inverted_image_data = 1.0 - image_data
                    inverted_image_data = np.clip(inverted_image_data, 0.0, 1.0) # Ensure range for floats
                    logger.debug(f"Asset '{asset_name_for_log}': Inverted float image data for {original_temp_path}.")
                elif np.issubdtype(image_data.dtype, np.integer):
                    max_val = np.iinfo(image_data.dtype).max
                    inverted_image_data = max_val - image_data
                    logger.debug(f"Asset '{asset_name_for_log}': Inverted integer image data (max_val: {max_val}) for {original_temp_path}.")
                else:
                    logger.error(
                        f"Asset '{asset_name_for_log}': Unsupported image data type {image_data.dtype} "
                        f"for GLOSS map {original_temp_path}. Cannot invert. Adding original rule."
                    )
                    new_files_to_process.append(file_rule)
                    continue
                
                # Save New Temporary (Roughness) Map
                # Using original_temp_path.suffix ensures we keep the format (e.g., .png, .exr)
                # Ensure file_rule.map_type exists before using sanitize_filename
                map_type_for_filename = file_rule.map_type if hasattr(file_rule, 'map_type') else "unknownmaptype"
                new_temp_filename = f"rough_from_gloss_{sanitize_filename(map_type_for_filename)}_{file_rule.id.hex}{original_temp_path.suffix}"
                new_temp_path = context.engine_temp_dir / new_temp_filename
                
                save_success = ipu.save_image(new_temp_path, inverted_image_data)

                if save_success:
                    logger.info(
                        f"Asset '{asset_name_for_log}': Converted GLOSS map {original_temp_path} "
                        f"to ROUGHNESS map {new_temp_path}."
                    )
                    
                    # Assuming FileRule has model_copy method
                    modified_file_rule = file_rule.model_copy(deep=True) if hasattr(file_rule, 'model_copy') else file_rule
                    modified_file_rule.map_type = "ROUGHNESS" # Ensure map_type can be set
                    
                    # Update context.processed_maps_details for the original file_rule.id.hex
                    context.processed_maps_details[map_detail_key]['temp_processed_file'] = str(new_temp_path)
                    context.processed_maps_details[map_detail_key]['original_map_type_before_conversion'] = "GLOSS"
                    context.processed_maps_details[map_detail_key]['notes'] = "Converted from GLOSS by GlossToRoughConversionStage"

                    new_files_to_process.append(modified_file_rule)
                    processed_a_gloss_map = True
                else:
                    logger.error(
                        f"Asset '{asset_name_for_log}': Failed to save inverted ROUGHNESS map to {new_temp_path} "
                        f"for original GLOSS map (ID: {map_detail_key}). Adding original rule."
                    )
                    new_files_to_process.append(file_rule)
            else: # Not a gloss map
                new_files_to_process.append(file_rule)

        context.files_to_process = new_files_to_process

        if processed_a_gloss_map:
            logger.info(
                f"Asset '{asset_name_for_log}': Gloss to Roughness conversion stage successfully processed one or more maps and updated file list."
            )
        else:
            logger.debug(
                f"Asset '{asset_name_for_log}': No gloss maps were successfully converted in GlossToRoughConversionStage. "
                f"File list for next stage contains original non-gloss maps and any gloss maps that failed conversion."
            )

        return context