<?php

namespace AskTamMCP\MCP;

defined('ABSPATH') || exit;

/**
 * MCP Server — handles the SSE connection and JSON-RPC message dispatch.
 *
 * MCP protocol flow:
 *   1. Client connects to /sse  → server sends "endpoint" event with POST URL
 *   2. Client POSTs JSON-RPC messages to /message
 *   3. Server sends "message" SSE events back with JSON-RPC responses
 *
 * This implementation uses a stateless synchronous approach:
 *   - SSE endpoint performs initialize + tools/list inline
 *   - /message endpoint handles individual tool calls synchronously
 */
class Server {

    private Registry $registry;

    public function __construct() {
        $this->registry = new Registry();
    }

    public function registerRoutes(): void {
        // SSE stream endpoint — AI client connects here
        register_rest_route('asktam-mcp/v1', '/sse', [
            'methods'             => 'GET',
            'callback'            => [$this, 'handleSse'],
            'permission_callback' => '__return_true',
        ]);

        // Message endpoint — AI client POSTs JSON-RPC here
        register_rest_route('asktam-mcp/v1', '/message', [
            'methods'             => 'POST',
            'callback'            => [$this, 'handleMessage'],
            'permission_callback' => '__return_true',
        ]);

        // Simple info endpoint
        register_rest_route('asktam-mcp/v1', '/info', [
            'methods'             => 'GET',
            'callback'            => [$this, 'handleInfo'],
            'permission_callback' => '__return_true',
        ]);
    }

    /**
     * SSE endpoint — opens the connection and sends server capabilities.
     */
    public function handleSse(\WP_REST_Request $request): void {
        $auth = \AskTamMCP\Core\Auth::validate($request, \AskTamMCP\Core\Auth::SCOPE_READ);
        if (is_wp_error($auth)) {
            status_header($auth->get_error_data()['status'] ?? 401);
            header('Content-Type: application/json');
            echo wp_json_encode(['error' => $auth->get_error_message()]);
            exit;
        }

        // SSE headers
        header('Content-Type: text/event-stream');
        header('Cache-Control: no-cache');
        header('X-Accel-Buffering: no');
        header('Connection: keep-alive');
        ob_end_flush();

        $messageEndpoint = rest_url('asktam-mcp/v1/message');

        // Send endpoint event — tells client where to POST messages
        $this->sseEvent('endpoint', $messageEndpoint);

        // Send server capabilities (MCP initialize response)
        $initResponse = [
            'jsonrpc' => '2.0',
            'id'      => null,
            'result'  => [
                'protocolVersion' => '2024-11-05',
                'capabilities'    => [
                    'tools' => ['listChanged' => false],
                ],
                'serverInfo' => [
                    'name'    => 'advanced-mcp-plugin',
                    'version' => ASKTAM_MCP_VERSION,
                ],
            ],
        ];
        $this->sseEvent('message', wp_json_encode($initResponse));

        // Keep-alive ping loop (max 29 seconds to avoid WP timeout)
        $start = time();
        while ((time() - $start) < 25) {
            $this->sseEvent('ping', (string) time());
            if (connection_aborted()) {
                break;
            }
            sleep(5);
        }

        exit;
    }

    /**
     * Message endpoint — receives JSON-RPC, dispatches to tool handler, returns result.
     */
    public function handleMessage(\WP_REST_Request $request): \WP_REST_Response|\WP_Error {
        $auth = \AskTamMCP\Core\Auth::validate($request, \AskTamMCP\Core\Auth::SCOPE_READ);
        if (is_wp_error($auth)) {
            return $auth;
        }

        $body = $request->get_json_params();

        if (empty($body['method'])) {
            return $this->rpcError(null, -32600, 'Invalid Request: missing method');
        }

        $id     = $body['id'] ?? null;
        $method = $body['method'];
        $params = $body['params'] ?? [];

        switch ($method) {
            case 'initialize':
                return $this->rpcResult($id, [
                    'protocolVersion' => '2024-11-05',
                    'capabilities'    => ['tools' => ['listChanged' => false]],
                    'serverInfo'      => ['name' => 'advanced-mcp-plugin', 'version' => ASKTAM_MCP_VERSION],
                ]);

            case 'tools/list':
                return $this->rpcResult($id, [
                    'tools' => $this->registry->getToolDefinitions(),
                ]);

            case 'tools/call':
                return $this->handleToolCall($id, $params, $auth);

            case 'notifications/initialized':
            case 'notifications/cancelled':
                // Client notifications — acknowledge silently
                return new \WP_REST_Response(null, 204);

            default:
                return $this->rpcError($id, -32601, "Method not found: {$method}");
        }
    }

    /**
     * Info endpoint — public, returns plugin version + endpoint URL.
     */
    public function handleInfo(): \WP_REST_Response {
        return new \WP_REST_Response([
            'name'     => 'Advanced MCP Plugin by AskTam',
            'version'  => ASKTAM_MCP_VERSION,
            'endpoint' => rest_url('asktam-mcp/v1/sse'),
            'docs'     => 'https://asktam.net/mcp-plugin',
        ]);
    }

    // -------------------------------------------------------------------------
    // Private helpers
    // -------------------------------------------------------------------------

    private function handleToolCall(mixed $id, array $params, array $authRow): \WP_REST_Response {
        $toolName  = $params['name'] ?? '';
        $arguments = $params['arguments'] ?? [];

        if (empty($toolName)) {
            return $this->rpcError($id, -32602, 'Missing tool name');
        }

        $handler = $this->registry->getHandler($toolName);
        if (!$handler) {
            return $this->rpcError($id, -32601, "Unknown tool: {$toolName}");
        }

        // Check write/admin scope for mutating tools
        $tool = $this->registry->getTool($toolName);
        if (($tool['scope'] ?? 'read') !== 'read') {
            $scopeCheck = \AskTamMCP\Core\Auth::validate(
                new \WP_REST_Request(), // dummy — scopes already on $authRow
                $tool['scope']
            );
            $keyScopes = array_map('trim', explode(',', $authRow['scopes']));
            $required  = $tool['scope'];
            if (!in_array($required, $keyScopes, true) && !in_array('admin', $keyScopes, true)) {
                return $this->rpcError($id, -32000, "Insufficient scope. Required: {$required}");
            }
        }

        try {
            $result = $handler($arguments);
            return $this->rpcResult($id, [
                'content' => [
                    ['type' => 'text', 'text' => is_string($result) ? $result : wp_json_encode($result)],
                ],
            ]);
        } catch (\Throwable $e) {
            return $this->rpcError($id, -32000, $e->getMessage());
        }
    }

    private function rpcResult(mixed $id, mixed $result): \WP_REST_Response {
        return new \WP_REST_Response([
            'jsonrpc' => '2.0',
            'id'      => $id,
            'result'  => $result,
        ]);
    }

    private function rpcError(mixed $id, int $code, string $message): \WP_REST_Response {
        return new \WP_REST_Response([
            'jsonrpc' => '2.0',
            'id'      => $id,
            'error'   => ['code' => $code, 'message' => $message],
        ], 200); // MCP errors use HTTP 200 with JSON-RPC error body
    }

    private function sseEvent(string $event, string $data): void {
        echo "event: {$event}\n";
        echo "data: {$data}\n\n";
        flush();
    }
}