<?php
/**
 * Copyright 2026 Kacper Donat
 *
 * @author Kacper "Kadet" Donat <kacper@kadet.net>
 *
 * Full license available in separate LICENSE file
 */

/**
 * Eager collection transforms.
 *
 * Every function returns a closure that takes the data, so they compose with the pipe operator. The result is the
 * same kind of container as the input: an array gives an array (keys preserved, unless stated otherwise), `Ds\Seq`
 * gives `Ds\Seq`, `Ds\Map` gives `Ds\Map` and `Ds\Set` gives `Ds\Set`. Any other iterable (e.g. a generator) is
 * consumed and gives a `Ds\Seq` of the results, or a list array when php-ds is unavailable.
 *
 * Callbacks receive `($value, $key)`, also for `Ds\Map`.
 */

namespace Kadet\Functional\Collections;

use Kadet\Functional\Collections\Lazy as lazy;
use function Kadet\Functional\Internal\callback;
use function Kadet\Functional\Collections\Internal\entries;
use function Kadet\Functional\Collections\Internal\iterated;
use function Kadet\Functional\Collections\Internal\keyedLike;
use function Kadet\Functional\Collections\Internal\seen;
use function Kadet\Functional\Collections\Internal\listLike;
use function Kadet\Functional\Collections\Internal\seqOf;
use function Kadet\Functional\Collections\Internal\unpairs;
use function Kadet\Functional\Collections\Internal\valuesOf;
use function Kadet\Functional\compile;

/**
 * Transforms each value with `$fn($value, $key)`, keeping the keys.
 *
 * For `Ds\Set` this follows `Ds\Set::map`, so equal results collapse into one.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 * @template TResult = mixed
 *
 * @param callable(TValue, TKey): TResult $fn
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? array<TKey, TResult>
 *     : (TIn is \Ds\Map<mixed, mixed>
 *         ? \Ds\Map<TKey, TResult>
 *         : (TIn is \Ds\Set<mixed>
 *             ? \Ds\Set<TResult>
 *             : (TIn is \Ds\Seq<mixed> ? \Ds\Seq<TResult> : \Ds\Seq<TResult>|list<TResult>))))
 */
function map(callable $fn): \Closure
{
    $fn = callback($fn);

    return static function (iterable $input) use ($fn): iterable {
        if (is_array($input)) {
            $result = [];
            foreach ($input as $key => $value) {
                $result[$key] = $fn($value, $key);
            }

            return $result;
        }

        if ($input instanceof \Ds\Map) {
            return $input->map(static fn ($key, $value) => $fn($value, $key));
        }

        $result = [];
        foreach (iterated($input) as $key => $value) {
            $result[] = $fn($value, $key);
        }

        return listLike($input, $result);
    };
}

/**
 * Keeps the values for which `$fn($value, $key)` is truthy, or the truthy values when `$fn` is null.
 *
 * Keys are preserved for arrays and `Ds\Map`; `Ds\Seq` is reindexed.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 *
 * @param (callable(TValue, TKey): mixed)|null $fn
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? array<TKey, TValue>
 *     : (TIn is \Ds\Map<mixed, mixed>
 *         ? \Ds\Map<TKey, TValue>
 *         : (TIn is \Ds\Set<mixed>
 *             ? \Ds\Set<TValue>
 *             : (TIn is \Ds\Seq<mixed> ? \Ds\Seq<TValue> : \Ds\Seq<TValue>|list<TValue>))))
 */
function filter(?callable $fn = null): \Closure
{
    $fn = $fn === null ? static fn ($value) => $value : callback($fn);

    return static function (iterable $input) use ($fn): iterable {
        if ($input instanceof \Ds\Map) {
            return $input->filter(static fn ($key, $value) => $fn($value, $key));
        }

        $values = iterated($input);
        if (is_array($values)) {
            $values = array_filter($values, $fn, ARRAY_FILTER_USE_BOTH);

            return is_array($input) ? $values : listLike($input, $values);
        }

        $result = [];
        foreach ($values as $key => $value) {
            if ($fn($value, $key)) {
                $result[] = $value;
            }
        }

        return listLike($input, $result);
    };
}

/**
 * Drops the values for which `$fn($value, $key)` is truthy. The inverse of {@see filter()}.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 *
 * @param callable(TValue, TKey): mixed $fn
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? array<TKey, TValue>
 *     : (TIn is \Ds\Map<mixed, mixed>
 *         ? \Ds\Map<TKey, TValue>
 *         : (TIn is \Ds\Set<mixed>
 *             ? \Ds\Set<TValue>
 *             : (TIn is \Ds\Seq<mixed> ? \Ds\Seq<TValue> : \Ds\Seq<TValue>|list<TValue>))))
 */
function reject(callable $fn): \Closure
{
    $fn = callback($fn);

    return static function (iterable $input) use ($fn): iterable {
        if ($input instanceof \Ds\Map) {
            return $input->filter(static fn ($key, $value) => !$fn($value, $key));
        }

        $result = [];
        if (is_array($input)) {
            foreach ($input as $key => $value) {
                if (!$fn($value, $key)) {
                    $result[$key] = $value;
                }
            }

            return $result;
        }

        foreach (iterated($input) as $key => $value) {
            if (!$fn($value, $key)) {
                $result[] = $value;
            }
        }

        return listLike($input, $result);
    };
}

/**
 * Maps each value to an iterable with `$fn($value, $key)` and flattens the results one level into a list.
 *
 * Arrays give a list array, `Ds\Set` gives `Ds\Set`, anything else gives `Ds\Seq`.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 * @template TResult = mixed
 *
 * @param callable(TValue, TKey): iterable<mixed, TResult> $fn
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? list<TResult>
 *     : (TIn is \Ds\Set<mixed>
 *         ? \Ds\Set<TResult>
 *         : (TIn is \Ds\Seq<mixed>|\Ds\Map<mixed, mixed> ? \Ds\Seq<TResult> : \Ds\Seq<TResult>|list<TResult>)))
 */
function flatMap(callable $fn): \Closure
{
    $fn = callback($fn);

    return static function (iterable $input) use ($fn): iterable {
        $result = [];
        foreach (iterated($input) as $key => $value) {
            $items = $fn($value, $key);

            if (!is_iterable($items)) {
                throw new \UnexpectedValueException(sprintf(
                    'flatMap(): the callback must return an iterable, got %s.',
                    get_debug_type($items),
                ));
            }

            foreach ($items as $item) {
                $result[] = $item;
            }
        }

        return listLike($input, $result);
    };
}

/**
 * Drops the keys: an array gives a list array, `Ds\Map` gives `Ds\Seq`.
 *
 * @return \Closure<TValue>(iterable<mixed, TValue>): iterable<int, TValue>
 */
function values(): \Closure
{
    return static fn (iterable $input): iterable => listLike($input, valuesOf($input));
}

/**
 * Gives the keys as a list: a list array for an array, a `Ds\Seq` for anything else.
 *
 * @return \Closure<TKey>(iterable<TKey, mixed>): iterable<int, TKey>
 */
function keys(): \Closure
{
    return static function (iterable $input): iterable {
        if (is_array($input)) {
            return array_keys($input);
        }

        $keys = [];
        foreach (iterated($input) as $key => $_) {
            $keys[] = $key;
        }

        return seqOf($keys);
    };
}

/**
 * Removes duplicate values, keeping the first occurrence. Values are compared strictly, objects by identity.
 *
 * @return \Closure<TKey, TValue>(iterable<TKey, TValue>): iterable<TKey, TValue>
 */
function unique(): \Closure
{
    return uniqueBy(static fn ($value) => $value);
}

/**
 * Removes values whose `$selector($value, $key)` result was already seen, keeping the first occurrence.
 * Selected values are compared strictly, objects by identity.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 *
 * @param callable(TValue, TKey): mixed $selector
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? array<TKey, TValue>
 *     : (TIn is \Ds\Map<mixed, mixed>
 *         ? \Ds\Map<TKey, TValue>
 *         : (TIn is \Ds\Set<mixed>
 *             ? \Ds\Set<TValue>
 *             : (TIn is \Ds\Seq<mixed> ? \Ds\Seq<TValue> : \Ds\Seq<TValue>|list<TValue>))))
 */
function uniqueBy(callable $selector): \Closure
{
    $selector = callback($selector);

    return static function (iterable $input) use ($selector): iterable {
        $seen = seen();

        if ($input instanceof \Ds\Map) {
            return $input->filter(static fn ($key, $value) => !$seen($selector($value, $key)));
        }

        $result = [];
        if (is_array($input)) {
            foreach ($input as $key => $value) {
                if (!$seen($selector($value, $key))) {
                    $result[$key] = $value;
                }
            }

            return $result;
        }

        foreach (iterated($input) as $key => $value) {
            if (!$seen($selector($value, $key))) {
                $result[] = $value;
            }
        }

        return listLike($input, $result);
    };
}

/**
 * Sorts the values with `$comparator($a, $b)` (an int like `<=>`), or with `<=>` when it is null.
 *
 * The sort is stable and keeps the keys of arrays and `Ds\Map`. The input is never modified.
 *
 * @template TValue = mixed
 *
 * @param (callable(TValue, TValue): int)|null $comparator
 *
 * @return \Closure<TKey>(iterable<TKey, TValue>): iterable<TKey, TValue>
 */
function sort(?callable $comparator = null): \Closure
{
    // Comparators always take two values, so they are compiled without the ($value, $key) adaptation.
    $comparator = $comparator === null ? static fn ($a, $b) => $a <=> $b : compile($comparator);

    return static function (iterable $input) use ($comparator): iterable {
        $entries = entries($input);
        usort($entries, static fn ($a, $b) => $comparator($a[1], $b[1]));

        return keyedLike($input, unpairs($entries));
    };
}

/**
 * Sorts the values by `$selector($value, $key)`, compared with `<=>`, in ascending or descending order.
 *
 * The sort is stable (equal values keep their order in both directions) and keeps the keys of arrays and
 * `Ds\Map`. The selector is called once per value. The input is never modified.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 *
 * @param callable(TValue, TKey): mixed $selector
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? array<TKey, TValue>
 *     : (TIn is \Ds\Map<mixed, mixed>
 *         ? \Ds\Map<TKey, TValue>
 *         : (TIn is \Ds\Set<mixed>
 *             ? \Ds\Set<TValue>
 *             : (TIn is \Ds\Seq<mixed> ? \Ds\Seq<TValue> : \Ds\Seq<TValue>|list<TValue>))))
 */
function sortBy(callable $selector, bool $descending = false): \Closure
{
    $selector = callback($selector);

    return static function (iterable $input) use ($selector, $descending): iterable {
        $entries = [];
        foreach ($input as $key => $value) {
            $entries[] = [$key, $value, $selector($value, $key)];
        }

        $direction = $descending ? -1 : 1;
        usort($entries, static fn ($a, $b) => $direction * ($a[2] <=> $b[2]));

        return keyedLike($input, unpairs($entries));
    };
}

/**
 * Takes `$length` values (all remaining when null) starting at position `$offset`, like `array_slice()`.
 *
 * Negative `$offset` counts from the end, negative `$length` stops that many values before the end. Keys of arrays
 * and `Ds\Map` are preserved. With non-negative arguments a generator is consumed only as far as needed.
 *
 * @return \Closure<TKey, TValue>(iterable<TKey, TValue>): iterable<TKey, TValue>
 */
function slice(int $offset, ?int $length = null): \Closure
{
    $lazy = lazy\slice($offset, $length);

    return static function (iterable $input) use ($offset, $length, $lazy): iterable {
        if (is_array($input)) {
            return array_slice($input, $offset, $length, true);
        }

        if ($input instanceof \Ds\Seq || $input instanceof \Ds\Map || $input instanceof \Ds\Set) {
            return $length === null ? $input->slice($offset) : $input->slice($offset, $length);
        }

        return keyedLike($input, $lazy($input));
    };
}

/**
 * Takes the first `$n` values. See {@see slice()}.
 *
 * @return \Closure<TKey, TValue>(iterable<TKey, TValue>): iterable<TKey, TValue>
 */
function take(int $n): \Closure
{
    if ($n < 0) {
        throw new \InvalidArgumentException(sprintf('take(): count must be >= 0, got %d', $n));
    }

    return slice(0, $n);
}

/**
 * Skips the first `$n` values. See {@see slice()}.
 *
 * @return \Closure<TKey, TValue>(iterable<TKey, TValue>): iterable<TKey, TValue>
 */
function skip(int $n): \Closure
{
    if ($n < 0) {
        throw new \InvalidArgumentException(sprintf('skip(): count must be >= 0, got %d', $n));
    }

    return slice($n);
}

/**
 * Puts `$values` before the values of the input, as a list (keys are dropped).
 *
 * @template TValue = mixed
 *
 * @param TValue ...$values
 *
 * @return \Closure(iterable<mixed, TValue>): iterable<int, TValue>
 */
function prepend(mixed ...$values): \Closure
{
    $values = array_values($values);

    return static fn (iterable $input): iterable => listLike($input, array_merge($values, valuesOf($input)));
}

/**
 * Puts `$values` after the values of the input, as a list (keys are dropped).
 *
 * @template TValue = mixed
 *
 * @param TValue ...$values
 *
 * @return \Closure(iterable<mixed, TValue>): iterable<int, TValue>
 */
function append(mixed ...$values): \Closure
{
    return concat($values);
}

/**
 * Puts the values of `$other` after the values of the input, as a list (keys are dropped).
 *
 * @template TValue = mixed
 *
 * @param iterable<mixed, TValue> $other
 *
 * @return \Closure(iterable<mixed, TValue>): iterable<int, TValue>
 */
function concat(iterable $other): \Closure
{
    return static fn (iterable $input): iterable => listLike($input, array_merge(valuesOf($input), valuesOf($other)));
}

/**
 * Merges `$other` into the input by key, later keys win. Integer keys of arrays are kept, unlike `array_merge()`.
 *
 * Arrays and `Ds\Map` are merged by key; for `Ds\Seq`, `Ds\Set` and other iterables the values are appended.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 *
 * @param iterable<TKey, TValue> $other
 *
 * @return \Closure(iterable<TKey, TValue>): iterable<TKey, TValue>
 */
function merge(iterable $other): \Closure
{
    return static function (iterable $input) use ($other): iterable {
        if (is_array($input)) {
            foreach ($other as $key => $value) {
                if (!is_scalar($key) && $key !== null) {
                    throw new \InvalidArgumentException(sprintf(
                        'merge(): key of type %s cannot be an array key',
                        get_debug_type($key),
                    ));
                }

                $input[$key] = $value;
            }

            return $input;
        }

        if ($input instanceof \Ds\Map) {
            $result = $input->copy();
            $result->putAll($other);

            return $result;
        }

        return listLike($input, array_merge(valuesOf($input), valuesOf($other)));
    };
}

/**
 * Calls `$fn($value, $key)` for each value, for its side effects, and returns the input unchanged.
 *
 * The return value of `$fn` is ignored. Arrays and php-ds collections are returned as they are; any other
 * iterable is consumed and its values are returned as a `Ds\Seq` (or a list array without php-ds).
 *
 * @template TKey = mixed
 * @template TValue = mixed
 *
 * @param callable(TValue, TKey): mixed $fn
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? array<TKey, TValue>
 *     : (TIn is \Ds\Map<mixed, mixed>
 *         ? \Ds\Map<TKey, TValue>
 *         : (TIn is \Ds\Set<mixed>
 *             ? \Ds\Set<TValue>
 *             : (TIn is \Ds\Seq<mixed> ? \Ds\Seq<TValue> : \Ds\Seq<TValue>|list<TValue>))))
 */
function each(callable $fn): \Closure
{
    $fn = callback($fn);

    return static function (iterable $input) use ($fn): iterable {
        if (is_array($input) || $input instanceof \Ds\Seq || $input instanceof \Ds\Map || $input instanceof \Ds\Set) {
            foreach (iterated($input) as $key => $value) {
                $fn($value, $key);
            }

            return $input;
        }

        $values = [];
        foreach ($input as $key => $value) {
            $fn($value, $key);
            $values[] = $value;
        }

        return listLike($input, $values);
    };
}

/**
 * Alias of {@see each()}.
 *
 * @template TKey = mixed
 * @template TValue = mixed
 *
 * @param callable(TValue, TKey): mixed $fn
 *
 * @return \Closure<TIn of iterable<mixed, mixed>>(TIn): (TIn is array
 *     ? array<TKey, TValue>
 *     : (TIn is \Ds\Map<mixed, mixed>
 *         ? \Ds\Map<TKey, TValue>
 *         : (TIn is \Ds\Set<mixed>
 *             ? \Ds\Set<TValue>
 *             : (TIn is \Ds\Seq<mixed> ? \Ds\Seq<TValue> : \Ds\Seq<TValue>|list<TValue>))))
 */
function tap(callable $fn): \Closure
{
    return each($fn);
}