<?php
/**
 * This file is part of TypePHP.
 *
 * @link     https://www.swoole.com/
 * @contact  service@swoole.com
 */

namespace TypePhp\Generator;

use TypePhp\Type;

use PhpParser\Node;
use PhpParser\Node\Expr\Yield_;
use PhpParser\Node\Expr\YieldFrom;
use PhpParser\Node\Stmt\ClassMethod;
use PhpParser\Node\Stmt\Function_;
use PhpParser\Node\IntersectionType;
use PhpParser\Node\NullableType;
use PhpParser\Node\UnionType;
use TypePhp\Context\FunctionContext;
use TypePhp\Entity\FunctionDef;

trait FiberGenerator
{
    protected function containsYield(Function_|ClassMethod $v): bool
    {
        return $this->containsYieldInNodes($v->stmts ?? []);
    }

    protected function containsYieldInNodes(array $nodes): bool
    {
        foreach ($nodes as $node) {
            if ($node instanceof Node && $this->containsYieldInNode($node)) {
                return true;
            }
        }
        return false;
    }

    protected function containsYieldInNode(Node $node): bool
    {
        if ($node instanceof Yield_ || $node instanceof YieldFrom) {
            return true;
        }
        if ($node instanceof Node\FunctionLike || $node instanceof Node\Stmt\ClassLike) {
            return false;
        }
        foreach ($node->getSubNodeNames() as $name) {
            $subNode = $node->{$name};
            if ($subNode instanceof Node) {
                if ($this->containsYieldInNode($subNode)) {
                    return true;
                }
            } elseif (is_array($subNode) && $this->containsYieldInNodes($subNode)) {
                return true;
            }
        }
        return false;
    }

    protected function prepareGeneratorFunction(Function_|ClassMethod $v, FunctionDef $functionDef): void
    {
        if ($v->byRef) {
            $this->fatalError($v, 'Generators returning by reference are not supported yet');
        }
        foreach ($v->params as $param) {
            if ($param->byRef || $param->variadic) {
                $this->fatalError($param, 'Generators with by-reference or variadic parameters are not supported yet');
            }
        }
        if (!$this->generatorReturnTypeAcceptsFiber($v->returnType)) {
            $this->fatalError($v, 'Generator return type must accept \\FiberGenerator; use Iterator, Traversable, iterable, object, mixed, or omit the return type');
        }
        $functionDef->generator = true;
        $functionDef->returnType = Type::VAR;
        $functionDef->returnClass = '';
        $functionDef->returnTypeCheck = null;
        $functionDef->returnTypeStr = '';
        $functionDef->returnTypeNode = null;
    }

    protected function generatorReturnTypeAcceptsFiber(?Node $type): bool
    {
        if ($type === null) {
            return true;
        }
        if ($type instanceof NullableType) {
            return $this->generatorReturnTypeAcceptsFiber($type->type);
        }
        if ($type instanceof UnionType) {
            foreach ($type->types as $member) {
                if ($this->generatorReturnTypeAcceptsFiber($member)) {
                    return true;
                }
            }
            return false;
        }
        if ($type instanceof IntersectionType) {
            foreach ($type->types as $member) {
                if (!$this->generatorReturnTypeAcceptsFiber($member)) {
                    return false;
                }
            }
            return true;
        }

        $typeName = strtolower($this->parseIdentifier($type));
        if (in_array($typeName, ['mixed', 'object', 'iterable'], true)) {
            return true;
        }

        [, $class] = $this->resolveTypeDecl($type, self::DECL_TYPE_OF_RETURN);
        $class = strtolower(ltrim($class, '\\'));
        return in_array($class, ['iterator', 'traversable', 'fibergenerator'], true);
    }

    protected function parseYieldExpr(Yield_ $expr): string
    {
        if (!$this->inGeneratorBody) {
            $this->fatalError($expr, 'The `Expr_Yield` is not supported outside generator functions');
        }
        return 'typephp_fiber_suspend(' . $this->genYieldPayload($expr) . ', nullptr)';
    }

    protected function parseYieldStmt(Yield_ $expr): string
    {
        $payload = $this->genYieldPayload($expr);
        $closed = $this->genTmpVarName();
        $this->addLocalVar($closed, Type::BOOL);
        return $closed . ' = false;' . PHP_EOL
            . $this->getIndent() . $closed . ' = typephp_fiber_yield(' . $payload . ');' . PHP_EOL
            . $this->getIndent() . 'if (' . $closed . ') {' . PHP_EOL
            . $this->getIndent() . '    return ' . self::VALUE_NULL . ';' . PHP_EOL
            . $this->getIndent() . '}';
    }

    protected function parseYieldFromStmt(YieldFrom $expr): string
    {
        $closed = $this->genTmpVarName();
        $this->addLocalVar($closed, Type::BOOL);
        return $closed . ' = false;' . PHP_EOL
            . $this->getIndent() . 'typephp_fiber_yield_from(' . $this->genYieldFromIterable($expr) . ', &' . $closed . ');' . PHP_EOL
            . $this->getIndent() . 'if (' . $closed . ') {' . PHP_EOL
            . $this->getIndent() . '    return ' . self::VALUE_NULL . ';' . PHP_EOL
            . $this->getIndent() . '}';
    }

    protected function genYieldPayload(Yield_ $expr): string
    {
        if ($expr->key) {
            // PHP evaluates an explicit yield key before its value. Materialize
            // both operands so lowering helpers cannot hoist value side effects
            // ahead of the key or defer postfix side effects until after resume.
            $key = $this->materializeYieldOperand($expr->key, true);
            $value = $expr->value
                ? $this->materializeYieldOperand($expr->value, true)
                : self::VALUE_NULL;
            return 'php::Array(php::StdStrKeyMap{{"key", ' . $key . '}, {"value", ' . $value . '}, {"has_key", true}})';
        }
        $value = $expr->value
            ? $this->materializeYieldOperand($expr->value)
            : self::VALUE_NULL;
        return 'php::Array(php::StdStrKeyMap{{"value", ' . $value . '}, {"has_key", false}})';
    }

    private function materializeYieldOperand(Node $expr, bool $force = false): string
    {
        [$value, $beforeStmts, $afterStmts] = $this->parseExprWithCapturedStmts($expr);
        foreach ($beforeStmts as $stmt) {
            $this->context->beforeStmtLines[] = $stmt;
        }
        if (!$force && !$afterStmts) {
            return $value;
        }

        $tmpVar = $this->addTmpVar(Type::VAR);
        $this->context->beforeStmtLines[] = $tmpVar . ' = ' . $value . ';';
        foreach ($afterStmts as $stmt) {
            $this->context->beforeStmtLines[] = $stmt;
        }
        return $tmpVar;
    }

    private function genYieldFromIterable(YieldFrom $expr): string
    {
        return $this->materializeYieldOperand($expr->expr);
    }

    protected function parseYieldFromExpr(YieldFrom $expr): string
    {
        if (!$this->inGeneratorBody) {
            $this->fatalError($expr, 'The `Expr_YieldFrom` is not supported outside generator functions');
        }
        return 'typephp_fiber_yield_from(' . $this->genYieldFromIterable($expr) . ', nullptr)';
    }

    protected function genFiberGeneratorFunction(Function_|ClassMethod $v, FunctionDef $functionDef, string $nativeName): string
    {
        $entryContext = $this->context;
        $entryIndent = $this->indentLevel;
        $entryInGeneratorBody = $this->inGeneratorBody;

        try {
            return $this->doGenFiberGeneratorFunction($v, $functionDef, $nativeName);
        } finally {
            $this->context = $entryContext;
            $this->indentLevel = $entryIndent;
            $this->inGeneratorBody = $entryInGeneratorBody;
        }
    }

    private function doGenFiberGeneratorFunction(Function_|ClassMethod $v, FunctionDef $functionDef, string $nativeName): string
    {
        $functionDeclCode = Type::VAR . ' ' . self::PREFIX . $nativeName . '(';
        if ($this->class) {
            $functionDeclCode .= Type::OBJECT . ' &this_';
            if ($functionDef->params) {
                $functionDeclCode .= ', ';
            }
        }
        $functionDeclCode .= $functionDef->params . ')';

        $uses = [];
        foreach ($functionDef->argInfoList as $argInfo) {
            $uses[] = $argInfo->name;
        }

        $code = $functionDeclCode . ' {' . PHP_EOL;
        $this->indentLevel++;
        foreach ($functionDef->argInfoList as $i => $argInfo) {
            if (!empty($argInfo->typeCheck)) {
                $code .= $this->genUnionParamCheck($argInfo, $i);
            }
        }
        foreach ($functionDef->argInfoList as $argInfo) {
            if ($argInfo->property) {
                $code .= $this->getIndent() . $this->genPropertyPromotion($argInfo);
            }
        }

        $closureVar = $this->genTmpVarName();
        $code .= $this->getIndent() . 'php::ClosureFn ' . $closureVar . ' = []('
            . 'INTERNAL_FUNCTION_PARAMETERS, '
            . Type::OBJECT . ' &this_, '
            . Type::ARGS . ' &vars_) -> ' . Type::VAR . ' {' . PHP_EOL;

        $outerContext = $this->context;
        $outerIndent = $this->indentLevel;
        $outerInGeneratorBody = $this->inGeneratorBody;
        $this->context = new FunctionContext();
        $this->context->inClosure = true;
        $this->inGeneratorBody = true;
        $this->indentLevel++;

        foreach ($functionDef->argInfoList as $i => $argInfo) {
            $code .= $this->getIndent() . Type::VAR . ' ' . $argInfo->name . ' = vars_.get(' . $i . ');' . PHP_EOL;
            $this->addArgument($argInfo->name, Type::VAR);
        }
        if ($this->class) {
            $this->addArgument('this_', Type::OBJECT);
        }

        $body = '';
        $this->indentLevel++;
        if ($this->methodDef && $this->methodDef->hasDynamicCall) {
            $body .= $this->genScopeSwitchCode();
        }
        if ($v->stmts) {
            $body .= $this->parseStmts($v->stmts);
        }
        $body .= $this->getIndent() . 'return ' . self::VALUE_NULL . ';' . PHP_EOL;
        $this->indentLevel--;
        $code .= $this->genScopeVarDecl();
        $code .= $this->getIndent() . 'try {' . PHP_EOL;
        $code .= $body;
        $code .= $this->getIndent() . '} catch (zend_object *) {' . PHP_EOL;
        $code .= $this->getIndent() . '    return ' . self::VALUE_NULL . ';' . PHP_EOL;
        $code .= $this->getIndent() . '}' . PHP_EOL;

        $this->indentLevel = $outerIndent;
        $this->inGeneratorBody = $outerInGeneratorBody;
        $this->context = $outerContext;

        $code .= $this->getIndent() . '};' . PHP_EOL;
        $args = $uses ? '{ ' . implode(', ', $uses) . ' }' : '{}';
        $closureExpr = $this->class
            ? 'php::newClosure(' . $closureVar . ', ' . $args . ', this_)'
            : 'php::newClosure(' . $closureVar . ', ' . $args . ')';
        $code .= $this->getIndent() . 'return typephp_new_fiber_generator(' . $closureExpr . ');' . PHP_EOL;
        $this->indentLevel--;
        $code .= '}' . PHP_EOL;

        return $code;
    }
}