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

from .base_stage import ProcessingStage
from ..asset_context import AssetProcessingContext
from utils.path_utils import generate_path_from_pattern, sanitize_filename
from rule_structure import FileRule # Assuming these are needed for type hints if not directly in context


logger = logging.getLogger(__name__)

class OutputOrganizationStage(ProcessingStage):
    """
    Organizes output files by copying temporary processed files to their final destinations.
    """

    def execute(self, context: AssetProcessingContext) -> AssetProcessingContext:
        """
        Copies temporary processed and merged files to their final output locations
        based on path patterns and updates AssetProcessingContext.
        """
        asset_name_for_log = context.asset_rule.asset_name if hasattr(context, 'asset_rule') and context.asset_rule else "Unknown Asset"
        logger.debug(f"Asset '{asset_name_for_log}': Starting output organization stage.")

        if context.status_flags.get('skip_asset'):
            logger.info(f"Asset '{asset_name_for_log}': Output organization skipped as 'skip_asset' is True.")
            return context

        current_status = context.asset_metadata.get('status', '')
        if current_status.startswith("Failed") or current_status == "Skipped":
            logger.info(f"Asset '{asset_name_for_log}': Output organization skipped due to prior status: {current_status}.")
            return context

        final_output_files: List[str] = []
        overwrite_existing = False
        # Correctly access general_settings and overwrite_existing from config_obj
        if hasattr(context.config_obj, 'general_settings'):
            if isinstance(context.config_obj.general_settings, dict):
                overwrite_existing = context.config_obj.general_settings.get('overwrite_existing', False)
            elif hasattr(context.config_obj.general_settings, 'overwrite_existing'): # If general_settings is an object
                overwrite_existing = getattr(context.config_obj.general_settings, 'overwrite_existing', False)
        else:
            logger.warning(f"Asset '{asset_name_for_log}': config_obj.general_settings not found, defaulting overwrite_existing to False.")
        
        output_dir_pattern = getattr(context.config_obj, 'output_directory_pattern', "[supplier]/[assetname]")
        output_filename_pattern_config = getattr(context.config_obj, 'output_filename_pattern', "[assetname]_[maptype]_[resolution].[ext]")


        # A. Organize Processed Individual Maps
        if context.processed_maps_details:
            logger.debug(f"Asset '{asset_name_for_log}': Organizing {len(context.processed_maps_details)} processed individual map entries.")
            for processed_map_key, details in context.processed_maps_details.items():
                map_status = details.get('status')
                base_map_type = details.get('map_type', 'unknown_map_type') # Original map type

                if map_status in ['Processed', 'Processed_No_Variants']:
                    if not details.get('temp_processed_file'):
                        logger.debug(f"Asset '{asset_name_for_log}': Skipping map key '{processed_map_key}' (status '{map_status}') due to missing 'temp_processed_file'.")
                        details['status'] = 'Organization Skipped (Missing Temp File)'
                        continue
                    
                    temp_file_path = Path(details['temp_processed_file'])
                    resolution_str = details.get('processed_resolution_name', details.get('original_resolution_name', 'resX'))

                    token_data = {
                        "assetname": asset_name_for_log,
                        "supplier": context.effective_supplier or "DefaultSupplier",
                        "maptype": base_map_type,
                        "resolution": resolution_str,
                        "ext": temp_file_path.suffix.lstrip('.'),
                        "incrementingvalue": getattr(context, 'incrementing_value', None),
                        "sha5": getattr(context, 'sha5_value', None)
                    }
                    token_data_cleaned = {k: v for k, v in token_data.items() if v is not None}
                    
                    output_filename = generate_path_from_pattern(output_filename_pattern_config, token_data_cleaned)

                    try:
                        relative_dir_path_str = generate_path_from_pattern(
                            pattern_string=output_dir_pattern,
                            token_data=token_data_cleaned
                        )
                        final_path = Path(context.output_base_path) / Path(relative_dir_path_str) / Path(output_filename)
                        final_path.parent.mkdir(parents=True, exist_ok=True)

                        if final_path.exists() and not overwrite_existing:
                            logger.info(f"Asset '{asset_name_for_log}': Output file {final_path} for map '{processed_map_key}' exists and overwrite is disabled. Skipping copy.")
                        else:
                            shutil.copy2(temp_file_path, final_path)
                            logger.info(f"Asset '{asset_name_for_log}': Copied {temp_file_path} to {final_path} for map '{processed_map_key}'.")
                            final_output_files.append(str(final_path))
                        
                        details['final_output_path'] = str(final_path)
                        details['status'] = 'Organized'

                        # Update asset_metadata for metadata.json
                        map_metadata_entry = context.asset_metadata.setdefault('maps', {}).setdefault(processed_map_key, {})
                        map_metadata_entry['map_type'] = base_map_type
                        map_metadata_entry['path'] = str(Path(relative_dir_path_str) / Path(output_filename)) # Store relative path

                    except Exception as e:
                        logger.error(f"Asset '{asset_name_for_log}': Failed to copy {temp_file_path} for map key '{processed_map_key}'. Error: {e}", exc_info=True)
                        context.status_flags['output_organization_error'] = True
                        context.asset_metadata['status'] = "Failed (Output Organization Error)"
                        details['status'] = 'Organization Failed'

                elif map_status == 'Processed_With_Variants':
                    variants = details.get('variants')
                    if not variants: # No variants list, or it's empty
                        logger.warning(f"Asset '{asset_name_for_log}': Map key '{processed_map_key}' (status '{map_status}') has no 'variants' list or it is empty. Attempting fallback to base file.")
                        if not details.get('temp_processed_file'):
                            logger.error(f"Asset '{asset_name_for_log}': Skipping map key '{processed_map_key}' (fallback) as 'temp_processed_file' is also missing.")
                            details['status'] = 'Organization Failed (No Variants, No Temp File)'
                            continue # Skip to next map key

                        # Fallback: Process the base temp_processed_file
                        temp_file_path = Path(details['temp_processed_file'])
                        resolution_str = details.get('processed_resolution_name', details.get('original_resolution_name', 'baseRes'))

                        token_data = {
                            "assetname": asset_name_for_log,
                            "supplier": context.effective_supplier or "DefaultSupplier",
                            "maptype": base_map_type,
                            "resolution": resolution_str,
                            "ext": temp_file_path.suffix.lstrip('.'),
                            "incrementingvalue": getattr(context, 'incrementing_value', None),
                            "sha5": getattr(context, 'sha5_value', None)
                        }
                        token_data_cleaned = {k: v for k, v in token_data.items() if v is not None}
                        output_filename = generate_path_from_pattern(output_filename_pattern_config, token_data_cleaned)

                        try:
                            relative_dir_path_str = generate_path_from_pattern(
                                pattern_string=output_dir_pattern,
                                token_data=token_data_cleaned
                            )
                            final_path = Path(context.output_base_path) / Path(relative_dir_path_str) / Path(output_filename)
                            final_path.parent.mkdir(parents=True, exist_ok=True)

                            if final_path.exists() and not overwrite_existing:
                                logger.info(f"Asset '{asset_name_for_log}': Output file {final_path} for map '{processed_map_key}' (fallback) exists and overwrite is disabled. Skipping copy.")
                            else:
                                shutil.copy2(temp_file_path, final_path)
                                logger.info(f"Asset '{asset_name_for_log}': Copied {temp_file_path} to {final_path} for map '{processed_map_key}' (fallback).")
                                final_output_files.append(str(final_path))
                            
                            details['final_output_path'] = str(final_path)
                            details['status'] = 'Organized (Base File Fallback)'

                            map_metadata_entry = context.asset_metadata.setdefault('maps', {}).setdefault(processed_map_key, {})
                            map_metadata_entry['map_type'] = base_map_type
                            map_metadata_entry['path'] = str(Path(relative_dir_path_str) / Path(output_filename))
                            if 'variant_paths' in map_metadata_entry: # Clean up if it was somehow set
                                del map_metadata_entry['variant_paths']
                        except Exception as e:
                            logger.error(f"Asset '{asset_name_for_log}': Failed to copy {temp_file_path} (fallback) for map key '{processed_map_key}'. Error: {e}", exc_info=True)
                            context.status_flags['output_organization_error'] = True
                            context.asset_metadata['status'] = "Failed (Output Organization Error - Fallback)"
                            details['status'] = 'Organization Failed (Fallback)'
                        continue # Finished with this map key due to fallback

                    # If we are here, 'variants' list exists and is not empty. Proceed with variant processing.
                    logger.debug(f"Asset '{asset_name_for_log}': Organizing {len(variants)} variants for map key '{processed_map_key}' (map type: {base_map_type}).")
                    
                    map_metadata_entry = context.asset_metadata.setdefault('maps', {}).setdefault(processed_map_key, {})
                    map_metadata_entry['map_type'] = base_map_type
                    map_metadata_entry.setdefault('variant_paths', {}) # Initialize if not present

                    processed_any_variant_successfully = False
                    failed_any_variant = False

                    for variant_index, variant_detail in enumerate(variants):
                        temp_variant_path_str = variant_detail.get('temp_path')
                        if not temp_variant_path_str:
                            logger.warning(f"Asset '{asset_name_for_log}': Variant {variant_index} for map '{processed_map_key}' is missing 'temp_path'. Skipping.")
                            variant_detail['status'] = 'Organization Skipped (Missing Temp Path)'
                            continue
                        
                        temp_variant_path = Path(temp_variant_path_str)
                        variant_resolution_key = variant_detail.get('resolution_key', f"varRes{variant_index}")
                        variant_ext = temp_variant_path.suffix.lstrip('.')

                        token_data_variant = {
                            "assetname": asset_name_for_log,
                            "supplier": context.effective_supplier or "DefaultSupplier",
                            "maptype": base_map_type,
                            "resolution": variant_resolution_key,
                            "ext": variant_ext,
                            "incrementingvalue": getattr(context, 'incrementing_value', None),
                            "sha5": getattr(context, 'sha5_value', None)
                        }
                        token_data_variant_cleaned = {k: v for k, v in token_data_variant.items() if v is not None}
                        output_filename_variant = generate_path_from_pattern(output_filename_pattern_config, token_data_variant_cleaned)
                        
                        try:
                            relative_dir_path_str_variant = generate_path_from_pattern(
                                pattern_string=output_dir_pattern,
                                token_data=token_data_variant_cleaned
                            )
                            final_variant_path = Path(context.output_base_path) / Path(relative_dir_path_str_variant) / Path(output_filename_variant)
                            final_variant_path.parent.mkdir(parents=True, exist_ok=True)

                            if final_variant_path.exists() and not overwrite_existing:
                                logger.info(f"Asset '{asset_name_for_log}': Output variant file {final_variant_path} for map '{processed_map_key}' (res: {variant_resolution_key}) exists and overwrite is disabled. Skipping copy.")
                                variant_detail['status'] = 'Organized (Exists, Skipped Copy)'
                            else:
                                shutil.copy2(temp_variant_path, final_variant_path)
                                logger.info(f"Asset '{asset_name_for_log}': Copied variant {temp_variant_path} to {final_variant_path} for map '{processed_map_key}'.")
                                final_output_files.append(str(final_variant_path))
                                variant_detail['status'] = 'Organized'
                            
                            variant_detail['final_output_path'] = str(final_variant_path)
                            relative_final_variant_path_str = str(Path(relative_dir_path_str_variant) / Path(output_filename_variant))
                            map_metadata_entry['variant_paths'][variant_resolution_key] = relative_final_variant_path_str
                            processed_any_variant_successfully = True

                        except Exception as e:
                            logger.error(f"Asset '{asset_name_for_log}': Failed to copy variant {temp_variant_path} for map key '{processed_map_key}' (res: {variant_resolution_key}). Error: {e}", exc_info=True)
                            context.status_flags['output_organization_error'] = True
                            context.asset_metadata['status'] = "Failed (Output Organization Error - Variant)"
                            variant_detail['status'] = 'Organization Failed'
                            failed_any_variant = True
                    
                    # Update parent map detail status based on variant outcomes
                    if failed_any_variant:
                        details['status'] = 'Organization Failed (Variants)'
                    elif processed_any_variant_successfully:
                        # Check if all processable variants were organized
                        all_attempted_organized = True
                        for v_detail in variants:
                            if v_detail.get('temp_path') and not v_detail.get('status', '').startswith('Organized'):
                                all_attempted_organized = False
                                break
                        if all_attempted_organized:
                             details['status'] = 'Organized (All Attempted Variants)'
                        else:
                             details['status'] = 'Partially Organized (Variants)'
                    elif not any(v.get('temp_path') for v in variants): # No variants had temp_paths to begin with
                        details['status'] = 'Processed_With_Variants (No Valid Variants to Organize)'
                    else: # Variants list existed, items had temp_paths, but none were successfully organized (e.g., all skipped due to existing file and no overwrite)
                        details['status'] = 'Organization Skipped (No Variants Copied/Needed)'


                else: # Other statuses like 'Skipped', 'Failed', 'Organization Failed' etc.
                    logger.debug(f"Asset '{asset_name_for_log}': Skipping map key '{processed_map_key}' (status: '{map_status}') for organization as it's not 'Processed', 'Processed_No_Variants', or 'Processed_With_Variants'.")
                    continue
        else:
            logger.debug(f"Asset '{asset_name_for_log}': No processed individual maps to organize.")

        # B. Organize Merged Maps
        if context.merged_maps_details:
            logger.debug(f"Asset '{asset_name_for_log}': Organizing {len(context.merged_maps_details)} merged map(s).")
            for merge_op_id, details in context.merged_maps_details.items(): # Use merge_op_id
                if details.get('status') != 'Processed' or not details.get('temp_merged_file'):
                    logger.debug(f"Asset '{asset_name_for_log}': Skipping merge op id '{merge_op_id}' due to status '{details.get('status')}' or missing temp file.")
                    continue

                temp_file_path = Path(details['temp_merged_file'])
                map_type = details.get('map_type', 'unknown_merged_map') # This is the output_map_type of the merge rule
                # Merged maps might not have a simple 'resolution' token like individual maps.
                # We'll use a placeholder or derive if possible.
                resolution_str = details.get('merged_resolution_name', 'mergedRes')


                token_data_merged = {
                    "assetname": asset_name_for_log,
                    "supplier": context.effective_supplier or "DefaultSupplier",
                    "maptype": map_type,
                    "resolution": resolution_str,
                    "ext": temp_file_path.suffix.lstrip('.'),
                    "incrementingvalue": getattr(context, 'incrementing_value', None),
                    "sha5": getattr(context, 'sha5_value', None)
                }
                token_data_merged_cleaned = {k: v for k, v in token_data_merged.items() if v is not None}

                output_filename_merged = generate_path_from_pattern(output_filename_pattern_config, token_data_merged_cleaned)
                
                try:
                    relative_dir_path_str_merged = generate_path_from_pattern(
                        pattern_string=output_dir_pattern,
                        token_data=token_data_merged_cleaned
                    )
                    final_path_merged = Path(context.output_base_path) / Path(relative_dir_path_str_merged) / Path(output_filename_merged)
                    final_path_merged.parent.mkdir(parents=True, exist_ok=True)

                    if final_path_merged.exists() and not overwrite_existing:
                        logger.info(f"Asset '{asset_name_for_log}': Output file {final_path_merged} exists and overwrite is disabled. Skipping copy for merged map.")
                    else:
                        shutil.copy2(temp_file_path, final_path_merged)
                        logger.info(f"Asset '{asset_name_for_log}': Copied merged map {temp_file_path} to {final_path_merged}")
                        final_output_files.append(str(final_path_merged))
                    
                    context.merged_maps_details[merge_op_id]['final_output_path'] = str(final_path_merged)
                    context.merged_maps_details[merge_op_id]['status'] = 'Organized'

                except Exception as e:
                    logger.error(f"Asset '{asset_name_for_log}': Failed to copy merged map {temp_file_path} to destination for merge op id '{merge_op_id}'. Error: {e}", exc_info=True)
                    context.status_flags['output_organization_error'] = True
                    context.asset_metadata['status'] = "Failed (Output Organization Error)"
                    context.merged_maps_details[merge_op_id]['status'] = 'Organization Failed'
        else:
            logger.debug(f"Asset '{asset_name_for_log}': No merged maps to organize.")

        # C. Organize Extra Files (e.g., previews, text files)
        logger.debug(f"Asset '{asset_name_for_log}': Checking for EXTRA files to organize.")
        extra_files_organized_count = 0
        if hasattr(context, 'files_to_process') and context.files_to_process:
            extra_subdir_name = getattr(context.config_obj, 'extra_files_subdir', 'Extra') # Default to 'Extra'

            for file_rule in context.files_to_process:
                if file_rule.item_type == 'EXTRA':
                    source_file_path = context.workspace_path / file_rule.file_path
                    if not source_file_path.is_file():
                        logger.warning(f"Asset '{asset_name_for_log}': EXTRA file '{source_file_path}' not found. Skipping.")
                        continue

                    # Basic token data for the asset's base output directory
                    # We don't use map_type, resolution, or ext for the base directory of extras.
                    # However, generate_path_from_pattern might expect them or handle their absence.
                    # For the base asset directory, only assetname and supplier are typically primary.
                    base_token_data = {
                        "assetname": asset_name_for_log,
                        "supplier": context.effective_supplier or "DefaultSupplier",
                        # Add other tokens if your output_directory_pattern uses them at the asset level
                        "incrementingvalue": getattr(context, 'incrementing_value', None),
                        "sha5": getattr(context, 'sha5_value', None)
                    }
                    base_token_data_cleaned = {k: v for k, v in base_token_data.items() if v is not None}

                    try:
                        asset_base_output_dir_str = generate_path_from_pattern(
                            pattern_string=output_dir_pattern, # Uses the same pattern as other maps for base dir
                            token_data=base_token_data_cleaned
                        )
                        # Destination: <output_base_path>/<asset_base_output_dir_str>/<extra_subdir_name>/<original_filename>
                        final_dest_path = (Path(context.output_base_path) /
                                           Path(asset_base_output_dir_str) /
                                           Path(extra_subdir_name) /
                                           source_file_path.name) # Use original filename

                        final_dest_path.parent.mkdir(parents=True, exist_ok=True)

                        if final_dest_path.exists() and not overwrite_existing:
                            logger.info(f"Asset '{asset_name_for_log}': EXTRA file destination {final_dest_path} exists and overwrite is disabled. Skipping copy.")
                        else:
                            shutil.copy2(source_file_path, final_dest_path)
                            logger.info(f"Asset '{asset_name_for_log}': Copied EXTRA file {source_file_path} to {final_dest_path}")
                            final_output_files.append(str(final_dest_path))
                            extra_files_organized_count += 1
                        
                        # Optionally, add more detailed tracking for extra files in context.asset_metadata
                        # For example:
                        # if 'extra_files_details' not in context.asset_metadata:
                        #     context.asset_metadata['extra_files_details'] = []
                        # context.asset_metadata['extra_files_details'].append({
                        #     'source_path': str(source_file_path),
                        #     'destination_path': str(final_dest_path),
                        #     'status': 'Organized'
                        # })

                    except Exception as e:
                        logger.error(f"Asset '{asset_name_for_log}': Failed to copy EXTRA file {source_file_path} to destination. Error: {e}", exc_info=True)
                        context.status_flags['output_organization_error'] = True
                        context.asset_metadata['status'] = "Failed (Output Organization Error - Extra Files)"
                        # Optionally, update status for the specific file_rule if tracked
        
        if extra_files_organized_count > 0:
            logger.info(f"Asset '{asset_name_for_log}': Successfully organized {extra_files_organized_count} EXTRA file(s).")
        else:
            logger.debug(f"Asset '{asset_name_for_log}': No EXTRA files were processed or found to organize.")


        context.asset_metadata['final_output_files'] = final_output_files

        if context.status_flags.get('output_organization_error'):
            logger.error(f"Asset '{asset_name_for_log}': Output organization encountered errors. Status: {context.asset_metadata['status']}")
        else:
            logger.info(f"Asset '{asset_name_for_log}': Output organization complete. {len(final_output_files)} files placed.")
        
        logger.debug(f"Asset '{asset_name_for_log}': Output organization stage finished.")
        return context