"""Object mapping functionality."""

from typing import Callable, TypeVar

from pydantic import BaseModel

from mosh.pit import Pit

T = TypeVar("T", bound=BaseModel)
U = TypeVar("U", bound=BaseModel)


def map_object(obj: T, mapper_fn: Callable[[T], U]) -> U:
    """Map a BaseModel object to another using a function.

    Args:
        obj: Source BaseModel to map from.
        mapper_fn: Function that takes source object and returns target object.

    Returns:
        The mapped object returned by mapper_fn.

    Example:
        person = Person(name="Alice", email="alice@example.com", age=30)
        summary = map_object(person, lambda p: PersonSummary(name=p.name))
    """
    return mapper_fn(obj)


def map_pit(
    pit: Pit,
    mapper_fn: Callable[[BaseModel], BaseModel],
    target_pit: Pit | None = None,
) -> Pit:
    """Map all objects in a pit to a new/existing pit.

    Args:
        pit: Source pit containing objects to map.
        mapper_fn: Function to transform each object.
        target_pit: Optional target pit. If None, creates a new in-memory pit.

    Returns:
        The target pit containing mapped objects.

    Example:
        summaries_pit = map_pit(people_pit, lambda p: PersonSummary(name=p.name))
    """
    if target_pit is None:
        target_pit = Pit(f"{pit.name}_mapped")

    for stored in pit.all_stored():
        mapped_obj = mapper_fn(stored.obj)
        target_pit.put(mapped_obj)

    return target_pit


def filter_pit(
    pit: Pit,
    predicate: Callable[[BaseModel], bool],
    target_pit: Pit | None = None,
) -> Pit:
    """Filter objects in a pit based on a predicate.

    Args:
        pit: Source pit containing objects to filter.
        predicate: Function that returns True for objects to keep.
        target_pit: Optional target pit. If None, creates a new in-memory pit.

    Returns:
        The target pit containing filtered objects.

    Example:
        adults_pit = filter_pit(people_pit, lambda p: p.age >= 18)
    """
    if target_pit is None:
        target_pit = Pit(f"{pit.name}_filtered")

    for stored in pit.all_stored():
        if predicate(stored.obj):
            # Re-put with the same ID to preserve identity
            target_pit.put(stored.obj, id=stored.id, tags=stored.tags)

    return target_pit


def transform_pit(
    pit: Pit,
    mapper_fn: Callable[[BaseModel], BaseModel | None],
    target_pit: Pit | None = None,
) -> Pit:
    """Transform objects in a pit, optionally filtering.

    Like map_pit, but mapper_fn can return None to skip objects.

    Args:
        pit: Source pit containing objects to transform.
        mapper_fn: Function to transform each object. Return None to skip.
        target_pit: Optional target pit. If None, creates a new in-memory pit.

    Returns:
        The target pit containing transformed objects.
    """
    if target_pit is None:
        target_pit = Pit(f"{pit.name}_transformed")

    for stored in pit.all_stored():
        result = mapper_fn(stored.obj)
        if result is not None:
            target_pit.put(result)

    return target_pit
