use App\Models\User;
use Illuminate\Support\Collection;

/** @return Collection<int|string, Collection<int|string, Collection<int, User>>> */
function example(): Collection {
  return User::all()->groupBy(['id', 'name']);
}

use Illuminate\Support\Collection;
use PhpParser\Node\Expr\MethodCall;
use PHPStan\Analyser\Scope;
use PHPStan\Reflection\MethodReflection;
use PHPStan\Type\Constant\ConstantIntegerType;
use PHPStan\Type\DynamicMethodReturnTypeExtension;
use PHPStan\Type\Generic\GenericObjectType;
use PHPStan\Type\IntegerRangeType;
use PHPStan\Type\IntegerType;
use PHPStan\Type\StringType;
use PHPStan\Type\Type;
use PHPStan\Type\UnionType;

class CollectionGroupByReturnExtension implements DynamicMethodReturnTypeExtension {
    public function getClass(): string {
        return Collection::class;
    }

    public function isMethodSupported(MethodReflection $methodReflection): bool {
        return $methodReflection->getName() === 'groupBy';
    }

    public function getTypeFromMethodCall(
        MethodReflection $methodReflection,
        MethodCall $methodCall,
        Scope $scope
    ): ?Type {
        $fields = $methodCall->getArgs();
        if (\count($fields) < 1) {
            return null;
        }

        $groupBy = $fields[0]->value;
        $groupByType = $scope->getType($groupBy);

        $collectionType = $scope->getType($methodCall->var);

        if (
            // @phpstan-ignore-next-line
            ! $collectionType instanceof GenericObjectType ||
            ! $groupByType->isIterable()->yes()
        ) {
            return null;
        }

        $limit = $groupByType->getArraySize();
        if ($limit instanceof ConstantIntegerType) {
            $limit = $limit->getValue();
        } elseif ($limit instanceof IntegerRangeType) {
            $limit = $limit->getMin();
        }

        if (! \is_int($limit) || $limit < 2) {
            return null;
        }

        $inner = new GenericObjectType(
            Collection::class,
            $collectionType->getTypes(),
        );

        while ($limit > 0) {
            $inner = new GenericObjectType(
                Collection::class,
                [
                    new UnionType([new IntegerType(), new StringType()]),
                    $inner,
                ],
            );

            $limit -= 1;
        }

        return $inner;
    }
}