<?php declare(strict_types = 1);

namespace PHPStan\Type\Php;

use PhpParser\Node\Expr\FuncCall;
use PHPStan\Analyser\Scope;
use PHPStan\Reflection\FunctionReflection;
use PHPStan\Reflection\ParametersAcceptorSelector;
use PHPStan\Type\Constant\ConstantIntegerType;
use PHPStan\Type\Type;
use PHPStan\Type\TypeCombinator;
use PHPStan\Type\TypeUtils;

class CountFunctionReturnTypeExtension implements \PHPStan\Type\DynamicFunctionReturnTypeExtension
{

	public function isFunctionSupported(FunctionReflection $functionReflection): bool
	{
		return $functionReflection->getName() === 'count';
	}

	public function getTypeFromFunctionCall(
		FunctionReflection $functionReflection,
		FuncCall $functionCall,
		Scope $scope
	): Type
	{
		if (count($functionCall->args) < 1) {
			return ParametersAcceptorSelector::selectSingle($functionReflection->getVariants())->getReturnType();
		}

		if (count($functionCall->args) > 1) {
			$mode = $scope->getType($functionCall->args[1]->value);
			if ($mode->isSuperTypeOf(new ConstantIntegerType(\COUNT_RECURSIVE))->yes()) {
				return ParametersAcceptorSelector::selectSingle($functionReflection->getVariants())->getReturnType();
			}
		}

		$arrays = TypeUtils::getArrays($scope->getType($functionCall->args[0]->value));
		if (count($arrays) === 0) {
			return ParametersAcceptorSelector::selectSingle($functionReflection->getVariants())->getReturnType();
		}
		$countTypes = [];
		foreach ($arrays as $array) {
			$countTypes[] = $array->count();
		}

		return TypeCombinator::union(...$countTypes);
	}

}
