<?php

namespace AskTamMCP\MCP;

defined('ABSPATH') || exit;

/**
 * MCP Server — supports both transports:
 *
 * 1. Streamable HTTP (new, Claude.ai uses this):
 *    POST /sse  with Accept: application/json  → returns JSON-RPC response directly
 *    POST /sse  with Accept: text/event-stream → streams SSE response
 *
 * 2. Legacy SSE (old):
 *    GET  /sse  → opens SSE stream, sends endpoint event
 *    POST /message → handles JSON-RPC messages
 */
class Server {

    private Registry $registry;

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

    public function registerRoutes(): void {
        // Primary endpoint — accepts both GET (legacy SSE) and POST (streamable HTTP)
        register_rest_route('asktam-mcp/v1', '/sse', [
            [
                'methods'             => 'GET',
                'callback'            => [$this, 'handleSseGet'],
                'permission_callback' => '__return_true',
            ],
            [
                'methods'             => 'POST',
                'callback'            => [$this, 'handleSsePost'],
                'permission_callback' => '__return_true',
            ],
        ]);

        // Legacy message endpoint (used by old SSE transport)
        register_rest_route('asktam-mcp/v1', '/message', [
            'methods'             => 'POST',
            'callback'            => [$this, 'handleMessage'],
            'permission_callback' => '__return_true',
        ]);

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

    // -------------------------------------------------------------------------
    // Streamable HTTP transport — POST /sse
    // Claude.ai sends JSON-RPC here, we respond with JSON or SSE stream
    // -------------------------------------------------------------------------
    public function handleSsePost(\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;
        }

        $body   = $request->get_json_params() ?? [];
        $accept = $request->get_header('Accept') ?? '';

        // Determine if client wants SSE stream or plain JSON
        $wantsStream = str_contains($accept, 'text/event-stream');

        // Handle batch requests (array of JSON-RPC calls)
        if (isset($body[0])) {
            $results = [];
            foreach ($body as $msg) {
                $result = $this->dispatchJsonRpc($msg, $auth);
                if ($result !== null) {
                    $results[] = $result;
                }
            }
            if ($wantsStream) {
                $this->streamJsonRpcResponses($results);
            } else {
                $this->jsonResponse($results);
            }
            return;
        }

        // Single JSON-RPC message
        $result = $this->dispatchJsonRpc($body, $auth);

        if ($wantsStream) {
            $this->streamJsonRpcResponses($result ? [$result] : []);
        } else {
            $this->jsonResponse($result ?? new \stdClass());
        }
    }

    // -------------------------------------------------------------------------
    // Legacy SSE transport — GET /sse
    // -------------------------------------------------------------------------
    public function handleSseGet(\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;
        }

        @set_time_limit(0);
        @ini_set('max_execution_time', 0);
        while (ob_get_level()) ob_end_clean();

        header('Content-Type: text/event-stream; charset=UTF-8');
        header('Cache-Control: no-cache, no-store');
        header('X-Accel-Buffering: no');
        header('Connection: keep-alive');
        header('Access-Control-Allow-Origin: *');

        $messageEndpoint = rest_url('asktam-mcp/v1/message');
        $this->sseEvent('endpoint', $messageEndpoint);
        $this->sseEvent('message', wp_json_encode([
            'jsonrpc' => '2.0',
            'id'      => null,
            'result'  => [
                'protocolVersion' => '2024-11-05',
                'capabilities'    => ['tools' => ['listChanged' => false]],
                'serverInfo'      => ['name' => 'advanced-mcp-plugin', 'version' => ASKTAM_MCP_VERSION],
            ],
        ]));

        $start = time();
        while (!connection_aborted() && (time() - $start) < 270) {
            $this->sseEvent('ping', (string) time());
            for ($i = 0; $i < 15; $i++) {
                if (connection_aborted()) break 2;
                sleep(1);
            }
        }

        echo "retry: 1000\n\n";
        flush();
        exit;
    }

    // -------------------------------------------------------------------------
    // Legacy message handler — POST /message
    // -------------------------------------------------------------------------
    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();
        $result = $this->dispatchJsonRpc($body ?? [], $auth);

        return new \WP_REST_Response(
            $result ?? ['jsonrpc' => '2.0', 'id' => null, 'result' => null],
            200
        );
    }

    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'),
            'transport' => ['streamable-http', 'sse'],
        ]);
    }

    // -------------------------------------------------------------------------
    // Core JSON-RPC dispatcher — shared by all transports
    // -------------------------------------------------------------------------
    private function dispatchJsonRpc(array $body, array $authRow): ?array {
        if (empty($body['method'])) return null;

        $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, $authRow);

            case 'notifications/initialized':
            case 'notifications/cancelled':
                return null; // notifications get no response

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

    // -------------------------------------------------------------------------
    // Output helpers
    // -------------------------------------------------------------------------
    private function streamJsonRpcResponses(array $responses): void {
        while (ob_get_level()) ob_end_clean();
        header('Content-Type: text/event-stream; charset=UTF-8');
        header('Cache-Control: no-cache, no-store');
        header('X-Accel-Buffering: no');
        header('Access-Control-Allow-Origin: *');

        foreach ($responses as $response) {
            $this->sseEvent('message', wp_json_encode($response));
        }

        // Send done marker
        echo "event: done\ndata: \n\n";
        flush();
        exit;
    }

    private function jsonResponse(mixed $data): void {
        while (ob_get_level()) ob_end_clean();
        header('Content-Type: application/json; charset=UTF-8');
        header('Access-Control-Allow-Origin: *');
        echo wp_json_encode($data);
        flush();
        exit;
    }

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

    private function handleToolCall(mixed $id, array $params, array $authRow): array {
        $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}");

        $tool     = $this->registry->getTool($toolName);
        $required = $tool['scope'] ?? 'read';
        if ($required !== 'read') {
            $keyScopes = array_map('trim', explode(',', $authRow['scopes']));
            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): array {
        return ['jsonrpc' => '2.0', 'id' => $id, 'result' => $result];
    }

    private function rpcError(mixed $id, int $code, string $message): array {
        return ['jsonrpc' => '2.0', 'id' => $id, 'error' => ['code' => $code, 'message' => $message]];
    }
}