<?php

namespace WordpressModels\Rest;

use Closure;
use Laravel\SerializableClosure\Exceptions\PhpVersionNotSupportedException;
use Laravel\SerializableClosure\SerializableClosure;
use Psr\Log\LoggerInterface;
use Symfony\Component\Cache\Adapter\AdapterInterface;
use Symfony\Component\Cache\Adapter\FilesystemAdapter;

/**
 * Class CachableRestServer
 *
 * This extension of the WP_REST_SERVER class allows us to cache the callback references of the registered routes.
 * Caching routes is only possible if the callback is a static method, if the class is a singleton, or if the class is
 * directly instantiable. Due to this, we hold a couple of cache keys to check if the callback is a static method, or if
 * the class is a singleton. If this is the case, we can cache the callback reference. If not, we can't cache the callback
 * reference, and we have to register the route as usual.
 *
 * For callbacks, we use lazy-loaded functions, where we only initialize controller classes when the route is called.
 * This dramatically improves the performance of the REST API.
 *
 * @todo WP_REST_Posts_Controllers are not cached, as they require a contructor parameter with a post_type.
 *       We probably can mitigate this by checking the class on reference creation and extracting the post_type from the contrller.
 * @todo Woocommerce Stats (mostly 'wc-analytics' namespace) controllers are not cached, we still find some unexpected behaviours
 *
 */
class CachableRestServer extends \WP_REST_Server
{

    private AdapterInterface $cache;
    private bool $isCacheHit;
    private array $endpointsToCache = [];
    private array $instantiatedCachedControllers = [];
    private array $uncacheableControllers = [
        'Automattic\WooCommerce\Admin\API\Reports\GenericController',
    ];
    private int $start;
    private ?LoggerInterface $logger = null;

    public function __construct(private ?string $key = null)
    {
        parent::__construct();
        $this->start = microtime(true);

        // first try to get the cache key from the 'rest_cache_key' filter, if not set, use the environment hash
        $this->key ??= apply_filters('rest_cache_key', '') ?: get_wp_environment_hash();
        $this->cache = new FilesystemAdapter('', 0, WP_CONTENT_DIR . '/cache/' . $this->key . '/rest_route_cache');
        $this->logger = apply_filters('rest_cache_logger', null);

        $enabled = (bool)get_option('rest_route_cache_enabled', true);

        // check whether cache is hit
        $endpointsItem = $this->cache->getItem('endpoints');
        $this->isCacheHit = $endpointsItem->isHit();
        if ($this->isCacheHit && $enabled) {
            // if cache is hit, we have to reconstruct the endpoints
            $this->endpoints = array_merge($this->endpoints + $this->reconstructEndpoints($endpointsItem->get()));

            // set the uncacheable controllers
            $uncacheableItem = $this->cache->getItem('uncachable_controllers');
            $this->uncacheableControllers = array_merge(
                $this->uncacheableControllers,
                $uncacheableItem->isHit() ? $uncacheableItem->get() ?? [] : []);

            // add controller deregistration hooks
            add_action('deregister_rest_controllers', [$this, 'deregisterDefaultRestControllers'], 9);
            do_action('deregister_rest_controllers', $this->uncacheableControllers);
        } else {
            // if not cache hit, we have to register the cache persisting action at the end of server initialization
            add_action('rest_api_init', [$this, 'persistCache'], PHP_INT_MAX);
        }
        // add cache invalidation hook and cli command
        add_action('rest_route_cache_invalidate', [$this, 'invalidateRouteCache']);
        \WP_CLI::add_command('rest-route-cache invalidate', [$this, 'invalidateRouteCache']);

        // add debug information
        add_action('rest_api_init', [$this, 'logStartupTime'], PHP_INT_MAX);
    }

    public function cliInvalidateCommand()
    {
        if (!$this->isCacheHit) {
            \WP_CLI::success('No cache hit, nothing to invalidate');
        }

        $this->invalidateRouteCache();
        \WP_CLI::success('Cache invalidated');
    }

    /**
     * Helper method to identify if the current request is a REST request.
     * @return bool
     */
    public function isRestRequest()
    {
        $restPrefix = rest_get_url_prefix();

        return array_key_exists('rest_route', $_REQUEST)
            || str_contains($_SERVER['REQUEST_URI'], $restPrefix)
            || (defined('REST_REQUEST') && REST_REQUEST);
    }

    /**
     * Logs the startup time of the REST API.
     * @return void
     */
    public function logStartupTime()
    {
        $end = microtime(true);
        $time = $end - $this->start;
        $this->logger?->debug(
            "CachedRestServer Debug Information",
            ['data' => [
                'startup_time' => round($time, 2) . " seconds",
                'cache_hit' => $this->isCacheHit,
                'uncacheable_controllers' => $this->uncacheableControllers,
                'is_rest_request' => $this->isRestRequest(),
                'route' => $_SERVER['REQUEST_URI'],
            ]]
        );
    }

    /**
     * Unregister any cacheable callback from the WP_Filter 'rest_api_init' hook.
     *
     * This method is called by the 'deregister_rest_controllers' action.
     *
     * @param array $uncachableControllers
     * @return void
     */
    public function deregisterDefaultRestControllers(array $uncachableControllers)
    {
        global $wp_filter;
        // get all 'rest_api_init' callbacks
        $callbacks = $wp_filter['rest_api_init']->callbacks;
        // remove all 'rest_api_init' callbacks where the class is a WP_Rest_Controller
        $wp_filter['rest_api_init']->callbacks = array_map(fn(array $prioCallbacks) => array_filter($prioCallbacks, function ($hookCallbacks) use ($uncachableControllers) {
            if (!is_array($hookCallbacks)) {
                return true;
            }
            $obj = $hookCallbacks['function'];
            // filter the callback, may be any other than an array-form callback, or a non-cacheable controller
            return !is_array($obj) || !array_filter($uncachableControllers, fn($c) => is_a($obj[0], $c));
        }), $callbacks);
    }

    /**
     * Invalidate the cache.
     * @return void
     */
    public function invalidateRouteCache()
    {
        $this->cache->clear();
    }

    /**
     * Registers a route to the server and stores any callback reference in the cache if possible.
     *
     * @inheritDoc
     */
    public function register_route($route_namespace, $route, $route_args, $override = false)
    {
        if ($this->isCacheHit && isset($this->endpoints[$route]) && !$override) {
            return;
        }

        if (!$this->isCacheHit) {
            foreach ($route_args as $key => $args) {
                if (is_array($args)) {
                    $arg = $this->buildCallbackCache($args);
                    if (isset($arg['callback_ref'])) {
                        $arg['namespace'] = $route_namespace;
                        $this->endpointsToCache[$route][] = $arg;
                    }
                } elseif (is_callable($args) && $ref = $this->createCallbackReference($args)) {
                    $this->endpointsToCache[$route][$key] = $ref;
                } else {
                    $this->endpointsToCache[$route][$key] = $args;
                }
            }
        }

        parent::register_route($route_namespace, $route, $route_args, $override);
    }

    /**
     * Persist the built cache.
     *
     * We also have to persist the uncacheable controllers, so we can deregister them on the next request.
     *
     * @return void
     * @throws \Psr\Cache\InvalidArgumentException
     */
    public function persistCache()
    {
        $unchachableItem = $this->cache->getItem('uncachable_controllers');
        $unchachableItem->set(array_unique($this->uncacheableControllers));
        $this->cache->save($unchachableItem);

        $endpointItem = $this->cache->getItem('endpoints');
        $endpointItem->set($this->endpointsToCache);
        $this->cache->save($endpointItem);
    }

    /**
     * Reconstructs cached callbacks.
     *
     * When there is a cache hit, we have to reconstruct the callbacks. This is done by creating a lazy-load function
     * for the callback. This function will be called when the route is called.
     *
     * @param array $cachedEndpoints
     * @return array
     */
    private function reconstructEndpoints(array $cachedEndpoints): array
    {
        $this->instantiatedCachedControllers[get_class($this)] = $this;
        $this->namespaces = [];
        foreach ($cachedEndpoints as $routeKey => &$routes) {
            $this->namespaces[$routes['namespace']] ??= [];
            $this->namespaces[$routes['namespace']][$routeKey] = true;
            if (array_is_list($routes)) {
                foreach ($routes as &$route) {
                    $route = $this->reconstructRoute($route);
                }
            } else {
                $routes = $this->reconstructRoute($routes);
            }

        }

        return $cachedEndpoints;
    }

    /**
     * Reconstructs the route args.
     *
     * @param array $args
     * @return array
     */
    private function reconstructRouteArgs(array $args): array
    {
        foreach ($args as $key => $arg) {
            if (isset($arg['sanitize_callback_ref'])) {
                $args[$key]['sanitize_callback'] = $this->createLazyCallback($arg['sanitize_callback_ref'], $arg['sanitize_callback_context'] ?? []);
            }

            if (isset($arg['validate_callback_ref'])) {
                $args[$key]['validate_callback'] = $this->createLazyCallback($arg['validate_callback_ref'], $arg['validate_callback_context'] ?? []);
            }

            if (isset($arg['properties'])) {
                $args[$key]['properties'] = $this->reconstructRouteArgs($arg['properties']);
            }

            if (isset($arg['items']) && !is_string($arg['items'])) {
                $args[$key]['items'] = $this->reconstructRouteArgs($arg['items']);
            }

            if (isset($arg['items']['properties'])) {
                $args[$key]['items']['properties'] = $this->reconstructRouteArgs($arg['items']['properties']);
            }

            if (isset($arg['items']['allOf'])) {
                $args[$key]['items']['allOf'] = $this->reconstructRouteArgs($arg['items']['allOf']);
            }

            if (isset($arg['items']['anyOf'])) {
                $args[$key]['items']['anyOf'] = $this->reconstructRouteArgs($arg['items']['anyOf']);
            }

            if (isset($arg['items']['oneOf'])) {
                $args[$key]['items']['oneOf'] = $this->reconstructRouteArgs($arg['items']['oneOf']);
            }
        }
        return $args;
    }

    /**
     * Create a lazy-load function for a callback reference.
     *
     * For static references, we can call the callback directly.
     *
     * For non-static class method references, we have to check if the class is a singleton. If it is, we can use the instance method to get the instance. If not, we have to create a
     * new instance of the class.
     *
     * For serialized closures, we have to unserialize the closure and return the closure.
     *
     * @param array|string $callbackReference
     * @param array{static?: bool, singleton?: bool} $context
     * @return Closure
     */
    public function createLazyCallback(array|string $callbackReference, array $context): Closure
    {
        return function (...$args) use ($callbackReference, $context) {
            if (is_array($callbackReference) && !($context['static'] ?? false)) {
                // if the callback is a static method, we can call it directly
                [$class, $method] = $callbackReference;
                // set the controller instance if it is not set yet
                // if the controller is a singleton, we can use the instance method to get the instance
                if (!isset($this->instantiatedCachedControllers[$class])) {
                    $reflector = new \ReflectionClass($class);

                    if ($context['singleton'] ?? false) {
                        foreach (['instance', 'getInstance', 'get_instance'] as $instanceMethod) {
                            if ($reflector->hasMethod($instanceMethod) && $reflector->getMethod($instanceMethod)->isStatic() && $reflector->getMethod($instanceMethod)->isPublic()) {
                                $this->instantiatedCachedControllers[$class] = $class::$instanceMethod();
                                break;
                            }
                        }
                    } else {
                        $this->instantiatedCachedControllers[$class] = new $class();
                    }
                }

                // set the callback
                $callback = [$this->instantiatedCachedControllers[$class], $method];
            } elseif (is_string($callbackReference) && is_serialized($callbackReference)) {
                // unserialize the callback
                $callback = unserialize($callbackReference);

                // if the callback is a closure, we have to check if it is a serializable closure
                if ($callback instanceof SerializableClosure) {
                    $callback = $callback->getClosure();
                }
            } else {
                $callback = $callbackReference;
            }

            // call the function
            return call_user_func_array($callback, $args);
        };
    }

    /**
     * @param array $args
     * @return array
     * @throws PhpVersionNotSupportedException
     * @throws \ReflectionException
     */
    public function buildCallbackCache(array $args): array
    {
        // cache build step, key is callback key and value is required flag
        $keys = [
            'callback' => true,
            'permission_callback' => false
        ];

        $arguments = $args;
        foreach ($keys as $key => $required) {
            if (!isset($arguments[$key])) {
                // if the callback is not set, we can skip this step
                continue;
            }
            $arguments = $this->replaceCallbackWithReference($arguments, $key);
            // check if the reference is set, and the callback is required
            if (!isset($arguments["{$key}_ref"]) && $required) {
                // if the callback's first index is an object, add to uncacheable controllers list
                if (is_array($args[$key]) && is_object($args[$key][0])) {
                    $this->uncacheableControllers[] = get_class($args[$key][0]);
                }

                return [];
            }
        }

        if (isset($arguments['args'])) {
            $arguments['args'] = $this->replaceParameterCallbacksWithReferences($arguments['args']);
        }

        return $arguments;
    }

    /**
     * Replaces the callback with a reference.
     *
     * For static references, we can replace the callback with a reference. For non-static references, we have to check
     * if the class is a singleton. If it is, we can replace the callback with a reference. If not, we have to register
     * the route as usual.
     *
     * @param array $args The arguments of the route
     * @param string $key The key of the callback, e.g. 'callback' or 'permission_callback'
     * @return array
     * @throws \ReflectionException
     */
    public function replaceCallbackWithReference(array $args, string $key): array
    {
        if (!isset($args[$key])) {
            return $args;
        }
        $callback = $args[$key];

        $arguments = $args;
        if (is_string($callback) || (is_array($callback) && is_string($callback[0]))) {
            unset($arguments[$key]);
            $arguments["{$key}_ref"] = $this->createCallbackReference($callback);
        } else {
            // the callback is a class method
            $object = is_array($callback) ? $callback[0] : $callback;

            if (is_callable($object)) {
                // if the object is callable, we can replace the callback with a reference
                unset($arguments[$key]);
                $arguments["{$key}_ref"] = $this->createCallbackReference($object);
                $arguments["{$key}_context"] = [];
            } elseif (!array_filter($this->uncacheableControllers, fn($c) => is_a($object, $c))) {
                // if the class is not in the uncacheable controllers list, we can replace the callback with a reference
                $arguments = $this->createClassMethodReference($object, $arguments, $key, $callback);
            }
        }

        return $arguments;
    }

    /**
     * @param callable|null $value
     * @return array|string|null
     * @throws PhpVersionNotSupportedException
     */
    public function createCallbackReference(?callable $value): array|string|null
    {
        if (is_array($value)) {
            [$classOrObject, $method] = $value;
            // get the class of the callback
            $class = is_object($classOrObject) ? get_class($classOrObject) : $classOrObject;
            // get the method of the callback
            $callback_reference = [$class, $method];
        } elseif (is_string($value)) {
            $callback_reference = $value;
        } elseif ($value instanceof Closure) {
            $callback_reference = serialize(new SerializableClosure($value));
        } elseif (is_object($value)) {
            $callback_reference = serialize(new SerializableClosure(fn(...$args) => $value(...$args)));
        }
        return $callback_reference ?? null;
    }

    /**
     * Recursively replaces all parameter callbacks with references.
     *
     * Support for WP API schema:
     * - object
     * - array
     * - scalar properties
     * - allOf
     * - anyOf
     * - oneOf
     *
     * @param array $parameters
     * @return array
     * @throws PhpVersionNotSupportedException
     */
    public function replaceParameterCallbacksWithReferences(array|string $parameters): array|string
    {
        if (is_string($parameters)) {
            return is_serialized($parameters) ?
                $this->replaceParameterCallbacksWithReferences(unserialize($parameters)) :
                $parameters;
        }
        foreach ($parameters as $key => $routeParameter) {
            foreach (['sanitize_callback', 'validate_callback'] as $callbackKey) {
                if (isset($routeParameter[$callbackKey])) {
                    $parameters[$key] = $this->replaceCallbackWithReference($parameters[$key], $callbackKey);
                }
            }

            if (isset($routeParameter['properties'])) {
                $parameters[$key]['properties'] = $this->replaceParameterCallbacksWithReferences($routeParameter['properties']);
            } elseif (isset($routeParameter['items'])) {
                $parameters[$key]['items'] = $this->replaceParameterCallbacksWithReferences($routeParameter['items']);

                if (isset($routeParameter['items']['properties'])) {
                    $parameters[$key]['items']['properties'] = $this->replaceParameterCallbacksWithReferences($routeParameter['items']['properties']);
                }

                // allOf
                if (isset($routeParameter['items']['allOf'])) {
                    $parameters[$key]['items']['allOf'] = $this->replaceParameterCallbacksWithReferences($routeParameter['items']['allOf']);
                }

                // anyOf
                if (isset($routeParameter['items']['anyOf'])) {
                    $parameters[$key]['items']['anyOf'] = $this->replaceParameterCallbacksWithReferences($routeParameter['items']['anyOf']);
                }

                // oneOf
                if (isset($routeParameter['items']['oneOf'])) {
                    $parameters[$key]['items']['oneOf'] = $this->replaceParameterCallbacksWithReferences($routeParameter['items']['oneOf']);
                }
            }

        }
        return $parameters;
    }

    /**
     * @param array $route
     * @return array
     */
    public function reconstructRoute(array $route): array
    {
        // get the callback reference
        $route['callback'] = $this->createLazyCallback($route['callback_ref'], $route['callback_context'] ?? []);

        if (isset($route['permission_callback_ref'])) {
            // get the permission callback reference
            $route['permission_callback'] = $this->createLazyCallback($route['permission_callback_ref'], $route['permission_callback_context'] ?? []);
        }

        // get the args
        if (isset($route['args'])) {
            $route['args'] = $this->reconstructRouteArgs($route['args']);
        }
        return $route;
    }

    /**
     * Create a reference for an array-style callback.
     *
     * We first check whether the given object or class is a singleton. If it is, we mark it as directly instantiable.
     * If the class has not constructor, or all parameters are optional, we can mark it as directly instantiable.
     * Else, we skip this controller and mark it as uncacheable.
     *
     * @param string|object $object
     * @param array $arguments
     * @param string $key
     * @param mixed $callback
     * @return array
     * @throws PhpVersionNotSupportedException
     * @throws \ReflectionException
     */
    public function createClassMethodReference(string|object $object, array $arguments, string $key, mixed $callback): array
    {
        // check if the class is a singleton
        $reflector = new \ReflectionClass($object);

        // check if the class has a static instance method
        $isDirectlyInstantiable = $isSingleton = array_reduce(
            ['instance', 'getInstance', 'get_instance'],
            fn($isSingleton, $instanceMethod) => $isSingleton
                || ($reflector->hasMethod($instanceMethod)
                    && $reflector->getMethod($instanceMethod)->isStatic()
                    && $reflector->getMethod($instanceMethod)->isPublic()), false);


        if (!$isSingleton) {
            // if the class has a constructor, and all parameters are optional, we can use the instance method
            $constructorParams = $reflector->getConstructor()?->getParameters();
            $isDirectlyInstantiable = !$reflector->getConstructor()
                || !array_filter($constructorParams, fn(\ReflectionParameter $param) => !$param->isOptional());
        }

        // check if the class is directly instantiable
        if ($isDirectlyInstantiable) {
            unset($arguments[$key]);
            $arguments["{$key}_ref"] = $this->createCallbackReference($callback);
            $arguments["{$key}_context"] = [];
            if ($isSingleton) {
                $arguments["{$key}_context"]['singleton'] = true;
            }
        }
        return $arguments;
    }

}
