"""
Collection Utilities

Common collection/list/dict manipulation functions.
"""

from typing import Any, Callable, Dict, List, TypeVar, Iterator, Optional
from itertools import groupby as itertools_groupby
from functools import reduce

T = TypeVar("T")
K = TypeVar("K")


def deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
    """
    Deep merge two dictionaries.
    
    Args:
        base: Base dictionary
        override: Dictionary to merge (takes precedence)
        
    Returns:
        Merged dictionary
    """
    result = base.copy()
    
    for key, value in override.items():
        if (
            key in result
            and isinstance(result[key], dict)
            and isinstance(value, dict)
        ):
            result[key] = deep_merge(result[key], value)
        else:
            result[key] = value
    
    return result


def flatten(nested: List[Any], depth: int = -1) -> List[Any]:
    """
    Flatten nested lists.
    
    Args:
        nested: Nested list
        depth: Maximum depth to flatten (-1 for unlimited)
        
    Returns:
        Flattened list
    """
    result = []
    
    def _flatten(items: List[Any], current_depth: int) -> None:
        for item in items:
            if isinstance(item, list) and (depth == -1 or current_depth < depth):
                _flatten(item, current_depth + 1)
            else:
                result.append(item)
    
    _flatten(nested, 0)
    return result


def chunk(items: List[T], size: int) -> List[List[T]]:
    """
    Split list into chunks of specified size.
    
    Args:
        items: List to chunk
        size: Chunk size
        
    Returns:
        List of chunks
    """
    if size <= 0:
        raise ValueError("Chunk size must be positive")
    
    return [items[i : i + size] for i in range(0, len(items), size)]


def unique(items: List[T], key: Optional[Callable[[T], Any]] = None) -> List[T]:
    """
    Remove duplicates while preserving order.
    
    Args:
        items: List with potential duplicates
        key: Optional key function for comparison
        
    Returns:
        List with duplicates removed
    """
    seen = set()
    result = []
    
    for item in items:
        k = key(item) if key else item
        if k not in seen:
            seen.add(k)
            result.append(item)
    
    return result


def group_by(items: List[T], key: Callable[[T], K]) -> Dict[K, List[T]]:
    """
    Group items by key function.
    
    Args:
        items: Items to group
        key: Key function
        
    Returns:
        Dictionary of grouped items
    """
    result: Dict[K, List[T]] = {}
    
    for item in items:
        k = key(item)
        if k not in result:
            result[k] = []
        result[k].append(item)
    
    return result


def pluck(items: List[Dict[str, Any]], key: str) -> List[Any]:
    """
    Extract values for a key from list of dicts.
    
    Args:
        items: List of dictionaries
        key: Key to extract
        
    Returns:
        List of values
    """
    return [item.get(key) for item in items]


def find(items: List[T], predicate: Callable[[T], bool]) -> Optional[T]:
    """
    Find first item matching predicate.
    
    Args:
        items: List to search
        predicate: Match function
        
    Returns:
        First matching item or None
    """
    for item in items:
        if predicate(item):
            return item
    return None


def find_index(items: List[T], predicate: Callable[[T], bool]) -> int:
    """
    Find index of first item matching predicate.
    
    Args:
        items: List to search
        predicate: Match function
        
    Returns:
        Index of first matching item or -1
    """
    for i, item in enumerate(items):
        if predicate(item):
            return i
    return -1


def partition(
    items: List[T],
    predicate: Callable[[T], bool],
) -> tuple[List[T], List[T]]:
    """
    Partition list into two based on predicate.
    
    Args:
        items: List to partition
        predicate: Partition function
        
    Returns:
        Tuple of (matching, non-matching)
    """
    matching = []
    non_matching = []
    
    for item in items:
        if predicate(item):
            matching.append(item)
        else:
            non_matching.append(item)
    
    return matching, non_matching


def omit(d: Dict[str, Any], keys: List[str]) -> Dict[str, Any]:
    """
    Create dict without specified keys.
    
    Args:
        d: Source dictionary
        keys: Keys to omit
        
    Returns:
        New dictionary without specified keys
    """
    return {k: v for k, v in d.items() if k not in keys}


def pick(d: Dict[str, Any], keys: List[str]) -> Dict[str, Any]:
    """
    Create dict with only specified keys.
    
    Args:
        d: Source dictionary
        keys: Keys to pick
        
    Returns:
        New dictionary with only specified keys
    """
    return {k: v for k, v in d.items() if k in keys}