# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
#

"""
Data manipulation and processing utilities.

This module contains reusable functions for working with data structures:
- Dictionary operations
- Data filtering and searching
- Serialization and transformation
"""

import random
import string
from datetime import datetime
from decimal import Decimal
from enum import Enum
from typing import Any

import simplejson as json


def search_dict(input_dict: dict, search_for: str) -> str | None:
    """
    Search a dictionary one level deep for a given string.

    Args:
        input_dict: Dictionary to search through
        search_for: String to search for

    Returns:
        str: Dictionary key where search string was found, or None if not found
    """
    for k in input_dict:
        for v in input_dict[k]:
            if search_for in v:
                return k
    return None


def collect_keys(d: dict, levels: int = 3, current_level: int = 1) -> set[str]:
    """
    Collect keys up to a certain level from a nested dictionary.

    Args:
        d: The dictionary to process.
        levels: The maximum depth of keys to collect.
        current_level: The current depth in the recursion.

    Returns:
        Set of keys collected up to the specified level.
    """
    if current_level > levels or not isinstance(d, dict):
        return set()

    keys = set(d.keys())
    for value in d.values():
        if isinstance(value, dict):
            keys.update(collect_keys(value, levels, current_level + 1))
    return keys


def search_nested_keys(nested_dict: dict, search_string: str) -> bool:
    """
    Collect 1st, 2nd, and 3rd level keys from a nested dictionary and search for a string.

    Args:
        nested_dict: The dictionary to search.
        search_string: The string to search for in the keys.

    Returns:
        True if the string is found in the collected keys, otherwise False.
    """
    collected_keys = collect_keys(nested_dict, levels=3)
    return search_string in collected_keys


def looks_like_a_connection_id(account_id: int | str) -> bool:
    """
    Detects if a given input looks like an AWS account ID

    Args:
        account_id: The ID to check

    Returns:
        Boolean: True if the input looks like an AWS account ID (12 digits)

    Examples:
        >>> looks_like_a_connection_id('123456789012')
        True

        >>> looks_like_a_connection_id('012345678901')
        True

        >>> looks_like_a_connection_id(123456789012)
        True

        >>> looks_like_a_connection_id('aws_account_id')
        False
    """
    if len(str(account_id)) != 12:
        return False
    try:
        int(account_id)
    except ValueError:
        return False
    return True


def generate_random_string(length: int = 12) -> str:
    """
    Generate a random string of letters and digits.

    Args:
        length: Length of the string to generate

    Returns:
        Randomly generated string

    Examples:
        >>> len(generate_random_string(32))
        32
    """
    return "".join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(length))


def filter_dict(input_data: dict, value: Any, field: str, exact: bool = False, depth: int = 1) -> dict:
    """
    Filters a dictionary based on a search value for a specified field at a given depth.
    If no field is specified, searches dictionary keys at the specified depth.

    Args:
        input_data: The dictionary to filter
        value: The value to search for
        field: The dictionary key to search in.
        exact: If True, perform a case-insensitive exact match. If False, perform substring match
        depth: The depth level to search within the nested dictionary (default: 1)

    Returns:
        dict: Filtered dictionary containing only matching key-value pairs

    Examples:
        >>> # Test exact match in field at depth 2
        >>> test_dict = {"store1": {"item1": {"name": "Apple"}, "item2": {"name": "Apple Pie"}}}
        >>> filter_dict(test_dict, "Apple", field="name", exact=True, depth=3)
        {'store1': {'item1': {'name': 'Apple'}}}

        >>> # Test with non-string value at depth 2
        >>> test_dict = {"store1": {"item1": {"price": 10}, "item2": {"price": 20}}}
        >>> filter_dict(test_dict, 10, field="price", depth=3)
        {'store1': {'item1': {'price': 10}}}

        >>> # Test with non-string value at depth 1
        >>> test_dict = {"item1": {"is_enabled": False}, "item2": {"is_enabled": True}}
        >>> filter_dict(test_dict, True, field="is_enabled", depth=2)
        {'item2': {'is_enabled': True}}
    """

    def safe_lower(val):
        # Convert value to lowercase string if it's a string, otherwise return as is
        return str(val).lower() if isinstance(val, str) else val

    def search_at_depth(d, val, fi, current_depth=1):
        if not isinstance(d, dict):
            return {}

        if current_depth == depth:
            if exact:
                # Search for exact match in specified field
                if safe_lower(d.get(fi)) == safe_lower(val):
                    return d
                else:
                    return None
            else:
                # Search for substring match in specified field
                if isinstance(val, str) and isinstance(d.get(fi), str) and safe_lower(val) in safe_lower(d.get(fi)) or val == d.get(fi):
                    return d
                else:
                    return None
        else:
            result = {}
            for k, v in d.items():
                if isinstance(v, dict):
                    nested_result = search_at_depth(v, val, fi, current_depth + 1)
                    if nested_result:
                        result[k] = nested_result
            return result

    return search_at_depth(input_data, value, field)


class JSONSanity(json.JSONEncoder):
    """
    Enhanced JSON encoder that handles common Python types.

    Handles sets, Decimal, Enum, and datetime types that standard
    JSON encoder cannot handle.
    """

    def default(self, o):
        if isinstance(o, set):
            return list(o)
        if isinstance(o, Decimal):
            return str(o)
        if isinstance(o, Enum):
            return o.value
        if isinstance(o, datetime):
            return o.isoformat()
        return super().default(o)


def chunks(lst: list, n: int) -> list[list]:
    """
    Split a list into chunks of size n.

    Args:
        lst: List to split
        n: Size of each chunk

    Returns:
        List of list chunks

    Examples:
        >>> chunks([1, 2, 3, 4, 5], 2)
        [[1, 2], [3, 4], [5]]
    """
    return [lst[i : i + n] for i in range(0, len(lst), n)]


def remove_nones(d: dict) -> dict:
    """
    Remove None values from a dictionary recursively.

    Args:
        d: Dictionary to process

    Returns:
        Dictionary with None values removed

    Examples:
        >>> remove_nones({'a': 1, 'b': None, 'c': {'d': None, 'e': 2}})
        {'a': 1, 'c': {'e': 2}}
    """
    if not isinstance(d, dict):
        return d

    result = {}
    for k, v in d.items():
        if v is None:
            continue
        if isinstance(v, dict):
            nested = remove_nones(v)
            if nested:  # Only add if not empty dict
                result[k] = nested
        elif isinstance(v, list):
            cleaned_list = []
            for item in v:
                if isinstance(item, dict):
                    cleaned_item = remove_nones(item)
                    if cleaned_item:  # Only add if not empty dict
                        cleaned_list.append(cleaned_item)
                elif item is not None:
                    cleaned_list.append(item)
            if cleaned_list:  # Only add if not empty list
                result[k] = cleaned_list
        else:
            result[k] = v
    return result
