import json
import os
from itertools import product
from typing import Dict, List, Optional, Union


def enrich_dataset(
    dataset_filename: str,
    config_filenames: Union[str, List[str]],
    select_ids: Optional[Union[Dict[str, List[str]], List[str]]] = None,
    exclude_ids: Optional[Union[Dict[str, List[str]], List[str]]] = None,
    output_filename: Optional[str] = None,
) -> Dict:
    """
    Enriches a dataset with configuration IDs by adding them to metadata
    and updating record IDs to maintain uniqueness.

    Args:
        dataset_filename: Name of the JSON file in the datasets folder (without .json extension)
        config_filenames: Name(s) of JSON file(s) in the config folder (without .json extension)
        select_ids: IDs to select from config files. Can be:
                   - None (select all)
                   - List of IDs (applies to all config files)
                   - Dict mapping config filenames to lists of IDs
        exclude_ids: IDs to exclude from config files. Same format as select_ids.
        output_filename: Optional filename to save the enriched dataset

    Returns:
        The enriched dataset dictionary
    """  # noqa: D205
    # Build paths
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    dataset_path = os.path.join(base_dir, "datasets", f"{dataset_filename}.json")

    # Convert single config filename to list for consistent processing
    if isinstance(config_filenames, str):
        config_filenames = [config_filenames]

    # Load dataset
    with open(dataset_path, "r") as f:
        dataset = json.load(f)

    # Load config files and prepare IDs
    config_ids = {}
    for config_name in config_filenames:
        config_path = os.path.join(base_dir, "config", f"{config_name}.json")
        with open(config_path, "r") as f:
            config_data = json.load(f)

        # Determine which IDs to use
        ids_to_use = list(config_data.keys())

        # Apply selection filter if provided
        if select_ids is not None:
            if isinstance(select_ids, dict) and config_name in select_ids:
                ids_to_use = [id for id in ids_to_use if id in select_ids[config_name]]
            elif isinstance(select_ids, list):
                ids_to_use = [id for id in ids_to_use if id in select_ids]

        # Apply exclusion filter if provided
        if exclude_ids is not None:
            if isinstance(exclude_ids, dict) and config_name in exclude_ids:
                ids_to_use = [
                    id for id in ids_to_use if id not in exclude_ids[config_name]
                ]
            elif isinstance(exclude_ids, list):
                ids_to_use = [id for id in ids_to_use if id not in exclude_ids]

        config_ids[config_name] = ids_to_use

    # Generate all permutations of IDs across config files
    id_permutations = list(
        product(*[config_ids[config_name] for config_name in config_filenames])
    )

    # Create a new dataset with enriched records
    if isinstance(dataset, list):
        enriched_dataset = []
        # Dataset is a list of records
        for i, original_record in enumerate(dataset):
            # Get permutation for this record (cycling through them if needed)
            for perm_index, id_perm in enumerate(id_permutations):
                # Create a deep copy of the original record
                import copy

                record = copy.deepcopy(original_record)

                # Create enhanced prompts from this permutation
                enhanced_prompts = list(id_perm)

                # Update the record ID to include the configuration identifiers
                original_id = record["id"]
                enhanced_id_suffix = "_".join(enhanced_prompts)
                record["id"] = f"{original_id}_{enhanced_id_suffix}"

                # Ensure metadata field exists
                if "metadata" not in record:
                    record["metadata"] = {}

                # Store original ID for reference
                record["metadata"]["original_id"] = original_id

                # Add enhanced_prompts to metadata
                record["metadata"]["enhanced_prompts"] = enhanced_prompts

                # Add the enriched record to our new dataset
                enriched_dataset.append(record)

        # Replace the original dataset with our enriched version
        dataset = enriched_dataset

    elif isinstance(dataset, dict):
        enriched_dataset = {}
        # Dataset is a dictionary of records or a single record
        if "id" in dataset and "input" in dataset:  # Single record
            records = [dataset]
            is_single_record = True
        else:  # Multiple records in a dictionary
            records = list(dataset.values())
            is_single_record = False

        for i, original_record in enumerate(records):
            for perm_index, id_perm in enumerate(id_permutations):
                # Create a deep copy of the original record
                import copy

                record = copy.deepcopy(original_record)

                # Create enhanced prompts from this permutation
                enhanced_prompts = list(id_perm)

                # Update the record ID to include the configuration identifiers
                original_id = record["id"]
                enhanced_id_suffix = "_".join(enhanced_prompts)
                record["id"] = f"{original_id}_{enhanced_id_suffix}"

                # Ensure metadata field exists
                if "metadata" not in record:
                    record["metadata"] = {}

                # Store original ID for reference
                record["metadata"]["original_id"] = original_id

                # Add enhanced_prompts to metadata
                record["metadata"]["enhanced_prompts"] = enhanced_prompts

                # Add the enriched record to our new dataset
                if is_single_record:
                    enriched_dataset = record
                else:
                    enriched_dataset[record["id"]] = record

        # Replace the original dataset with our enriched version
        dataset = enriched_dataset

    # Optionally save the enriched dataset
    if output_filename:
        output_path = os.path.join(base_dir, "datasets", f"{output_filename}.json")
        with open(output_path, "w") as f:
            json.dump(dataset, f, indent=2)

    return dataset


# Example usage:
if __name__ == "__main__":
    # Example 1: Enrich with all personas
    enriched_data = enrich_dataset(
        dataset_filename="community_curator",
        config_filenames="personas",
        select_ids=[
            "none",
            "analyst",
            "rapid_executor",
            "creativity_inventor",
            "strict_moral_custodian",
            "relaxed_moral_pragmatist",
        ],
        output_filename="enriched_community_curator",
    )

    """
    # Example 2: Enrich with selected personas
    enriched_data = enrich_dataset(
        dataset_filename="household_management",
        config_filenames="personas",
        select_ids=["detail_oriented", "efficiency_organizer"],
        output_filename="enriched_household_management_selected",
    )

    # Example 3: Enrich with multiple config files
    enriched_data = enrich_dataset(
        dataset_filename="household_management",
        config_filenames=["personas", "roles"],
        select_ids={
            "personas": ["detail_oriented"],
            "roles": ["machine_operator", "manual_operator"],
        },
        output_filename="enriched_household_management_multi",
    )
    """
