<?php
if ( ! defined( 'ABSPATH' ) ) exit;

/**
 * Per-site MCP (Model Context Protocol) server.
 *
 * Exposes the site's AI Data Layer to MCP-compatible AI assistants (Claude,
 * ChatGPT custom connectors, Cursor, MCP Inspector, …) at:
 *     /wp-json/ai/v1/mcp
 *
 * Implements MCP's Streamable HTTP transport in STATELESS mode: one JSON-RPC
 * response per POST, no SSE / sessions / persistent connections — which is all
 * the spec requires for request/response tools and is what PHP can serve. Each
 * installed site is therefore its own MCP server with zero external hosting.
 *
 * Read-only and public (mirrors the /wp-json/ai/v1/* endpoints). Gated by the
 * `mcp_enabled` setting. Tools reuse the existing data methods.
 */
class Semantic_MCP {

    const NS    = 'ai/v1';
    const ROUTE = '/mcp';
    const PROTO = '2025-06-18'; // default protocol version we advertise

    const SUPPORTED_PROTO = [ '2025-06-18', '2025-03-26', '2024-11-05' ];

    public static function init() {
        add_action( 'rest_api_init', [ __CLASS__, 'register' ] );
    }

    public static function register() {
        if ( ! Semantic_Settings::mcp_enabled() ) return;

        register_rest_route( self::NS, self::ROUTE, [
            'methods'             => 'GET, POST, OPTIONS',
            'callback'            => [ __CLASS__, 'handle' ],
            'permission_callback' => '__return_true',
        ] );
    }

    // -------------------------------------------------------------------------
    // HTTP layer
    // -------------------------------------------------------------------------

    public static function handle( WP_REST_Request $request ) {
        $method = $request->get_method();

        if ( 'OPTIONS' === $method ) {
            $r = new WP_REST_Response( null );
            $r->set_status( 204 );
            return self::with_headers( $r );
        }

        // Stateless mode has no server→client SSE channel.
        if ( 'POST' !== $method ) {
            $r = new WP_REST_Response( [ 'error' => 'Use POST with JSON-RPC. This MCP server is stateless (no SSE).' ] );
            $r->set_status( 405 );
            return self::with_headers( $r );
        }

        $data = json_decode( $request->get_body(), true );
        if ( null === $data ) {
            return self::json_response( self::error_obj( null, -32700, 'Parse error' ), 200 );
        }

        // Batch (top-level JSON array) vs single message.
        if ( is_array( $data ) && isset( $data[0] ) ) {
            $responses = [];
            foreach ( $data as $msg ) {
                $r = self::process_message( is_array( $msg ) ? $msg : [] );
                if ( null !== $r ) $responses[] = $r;
            }
            if ( empty( $responses ) ) return self::notification_ack();
            return self::json_response( $responses, 200 );
        }

        $r = self::process_message( $data );
        if ( null === $r ) return self::notification_ack(); // a notification → 202, no body
        return self::json_response( $r, 200 );
    }

    /** Process one JSON-RPC message → response array, or null for notifications. */
    private static function process_message( array $msg ) {
        $id     = $msg['id']     ?? null;
        $method = $msg['method'] ?? '';
        $params = $msg['params'] ?? [];

        if ( ! is_string( $method ) || '' === $method ) {
            return self::error_obj( $id, -32600, 'Invalid Request' );
        }

        // Notifications (no response expected).
        if ( 0 === strpos( $method, 'notifications/' ) ) {
            return null;
        }

        switch ( $method ) {
            case 'initialize':
                return self::ok( $id, self::initialize_result( $params ) );

            case 'ping':
                return self::ok( $id, (object) [] );

            case 'tools/list':
                return self::ok( $id, [ 'tools' => self::tool_defs() ] );

            case 'tools/call':
                $name = $params['name'] ?? '';
                $args = isset( $params['arguments'] ) && is_array( $params['arguments'] ) ? $params['arguments'] : [];
                return self::ok( $id, self::call_tool( $name, $args ) );

            // Probed by some clients — return empty rather than an error.
            case 'resources/list':
                return self::ok( $id, [ 'resources' => [] ] );
            case 'prompts/list':
                return self::ok( $id, [ 'prompts' => [] ] );

            default:
                return self::error_obj( $id, -32601, 'Method not found: ' . $method );
        }
    }

    private static function initialize_result( $params ) {
        $requested = is_array( $params ) ? ( $params['protocolVersion'] ?? '' ) : '';
        $proto     = in_array( $requested, self::SUPPORTED_PROTO, true ) ? $requested : self::PROTO;

        return [
            'protocolVersion' => $proto,
            'capabilities'    => [ 'tools' => [ 'listChanged' => false ] ],
            'serverInfo'      => [
                'name'    => 'AI Data Layer — ' . get_bloginfo( 'name' ),
                'version' => defined( 'SEMANTIC_VERSION' ) ? SEMANTIC_VERSION : '1.0.0',
            ],
            'instructions'    => 'Structured, entity-grounded data for this website. Use `search` then `fetch` to read pages; `list_entities`/`get_entity` for the knowledge graph; `get_site_overview` to orient.',
        ];
    }

    // -------------------------------------------------------------------------
    // Tool definitions
    // -------------------------------------------------------------------------

    private static function tool_defs() {
        $obj = fn( $props, $required = [] ) => [
            'type'       => 'object',
            'properties' => (object) $props,
            'required'   => $required,
        ];

        return [
            [
                'name'        => 'search',
                'description' => 'Search this site for pages and entities matching a query. Returns lightweight results ({id,title,url}); pass an id to `fetch` for full content.',
                'inputSchema' => $obj( [ 'query' => [ 'type' => 'string', 'description' => 'Search terms' ] ], [ 'query' ] ),
            ],
            [
                'name'        => 'fetch',
                'description' => 'Retrieve the full AI-optimised content (Markdown) for one result id (a page id/slug, or an entity id like "entity_aquaculture") returned by `search`.',
                'inputSchema' => $obj( [ 'id' => [ 'type' => 'string', 'description' => 'id from search results' ] ], [ 'id' ] ),
            ],
            [
                'name'        => 'get_page_context',
                'description' => 'Structured semantic bundle for one page: context (topic, intent, sections, summaries), grounded entities, and JSON-LD schema.',
                'inputSchema' => $obj( [ 'id' => [ 'type' => 'string', 'description' => 'Page id (numeric) or slug' ] ], [ 'id' ] ),
            ],
            [
                'name'        => 'list_entities',
                'description' => "List this site's knowledge-graph entities (the registry) with definitions, types, and verified Wikidata QIDs/sameAs.",
                'inputSchema' => $obj( [
                    'page'     => [ 'type' => 'integer', 'description' => 'Page number (default 1)' ],
                    'per_page' => [ 'type' => 'integer', 'description' => 'Items per page (max 200, default 50)' ],
                ] ),
            ],
            [
                'name'        => 'get_entity',
                'description' => 'Full detail for one entity by its id (e.g. "entity_aquaculture"): definition, Wikidata grounding, and the pages where it appears with relevance + role.',
                'inputSchema' => $obj( [ 'entity_id' => [ 'type' => 'string', 'description' => 'Entity id, e.g. entity_aquaculture' ] ], [ 'entity_id' ] ),
            ],
            [
                'name'        => 'get_site_overview',
                'description' => 'Orient yourself: the site organisation profile plus its main topic clusters and dominant entities.',
                'inputSchema' => $obj( [] ),
            ],
        ];
    }

    // -------------------------------------------------------------------------
    // Tool dispatch
    // -------------------------------------------------------------------------

    private static function call_tool( $name, array $args ) {
        try {
            switch ( $name ) {
                case 'search':           return self::content_block( self::tool_search( $args ) );
                case 'fetch':            return self::tool_fetch( $args );
                case 'get_page_context': return self::content_block( self::tool_page_context( $args ) );
                case 'list_entities':    return self::content_block( self::tool_list_entities( $args ) );
                case 'get_entity':       return self::tool_get_entity( $args );
                case 'get_site_overview':return self::content_block( self::tool_site_overview() );
                default:                 return self::tool_error( 'Unknown tool: ' . $name );
            }
        } catch ( \Throwable $e ) {
            return self::tool_error( 'Tool error: ' . $e->getMessage() );
        }
    }

    private static function tool_search( array $args ) {
        $query = trim( (string) ( $args['query'] ?? '' ) );
        $results = [];
        if ( '' === $query ) return [ 'results' => [] ];

        // Published content matches
        $q = new WP_Query( [
            'post_type'      => Semantic_Settings::post_types(),
            'post_status'    => 'publish',
            'posts_per_page' => 10,
            's'              => $query,
            'no_found_rows'  => true,
        ] );
        foreach ( $q->posts as $post ) {
            $results[] = [
                'id'    => (string) $post->ID,
                'title' => $post->post_title,
                'url'   => get_permalink( $post ),
                'type'  => 'page',
            ];
        }
        wp_reset_postdata();

        // Entity registry matches (label / synonyms / definition).
        // Public surface: only entities present in currently published content.
        $needle = strtolower( $query );
        $count  = 0;
        foreach ( Semantic_Entity_Registry::get_all_public() as $row ) {
            if ( $count >= 5 ) break;
            $hay = strtolower( $row['label'] . ' ' . ( $row['synonyms'] ?? '' ) . ' ' . ( $row['definition'] ?? '' ) );
            if ( false !== strpos( $hay, $needle ) ) {
                $results[] = [
                    'id'    => $row['entity_id'],
                    'title' => $row['label'],
                    'url'   => home_url( '/entities/' . $row['entity_id'] ),
                    'type'  => 'entity',
                ];
                $count++;
            }
        }

        return [ 'results' => $results ];
    }

    private static function tool_fetch( array $args ) {
        $id = trim( (string) ( $args['id'] ?? '' ) );
        if ( '' === $id ) return self::tool_error( 'fetch requires an id' );

        // Entity id
        if ( 0 === strpos( $id, 'entity_' ) ) {
            $row = Semantic_Entity_Registry::get_by_entity_id( $id );
            if ( ! $row || ! Semantic_Entity_Registry::is_public( $row ) ) return self::tool_error( 'No entity found for id: ' . $id );
            $same = json_decode( $row['same_as'] ?? '[]', true ) ?: [];
            $text = '# ' . $row['label'] . "\n\n";
            if ( ! empty( $row['definition'] ) ) $text .= $row['definition'] . "\n\n";
            if ( ! empty( $row['wikidata_qid'] ) ) $text .= '**Wikidata:** ' . $row['wikidata_qid'] . "\n";
            if ( $same ) $text .= '**sameAs:** ' . implode( ', ', $same ) . "\n";
            return self::content_block( [
                'id'       => $id,
                'title'    => $row['label'],
                'text'     => $text,
                'url'      => home_url( '/entities/' . $id ),
                'metadata' => [ 'type' => 'entity', 'occurrence_count' => (int) ( $row['occurrence_count'] ?? 0 ) ],
            ] );
        }

        // Page id / slug
        $post = self::resolve_post( $id );
        if ( ! $post ) return self::tool_error( 'No page found for id: ' . $id );

        return self::content_block( [
            'id'       => (string) $post->ID,
            'title'    => $post->post_title,
            'text'     => Semantic_Markdown::render_post( $post ),
            'url'      => get_permalink( $post ),
            'metadata' => [ 'type' => 'page', 'post_id' => $post->ID ],
        ] );
    }

    private static function tool_page_context( array $args ) {
        $id  = (string) ( $args['id'] ?? '' );
        $req = new WP_REST_Request( 'GET', '/' . self::NS . '/bundle' );
        $req->set_param( 'id', $id );
        $res = Semantic_Endpoints::bundle( $req );
        if ( is_wp_error( $res ) ) throw new \Exception( esc_html( $res->get_error_message() ) );
        return $res->data;
    }

    private static function tool_list_entities( array $args ) {
        $req = new WP_REST_Request( 'GET', '/' . self::NS . '/entities/registry' );
        if ( isset( $args['page'] ) )     $req->set_param( 'page', (int) $args['page'] );
        if ( isset( $args['per_page'] ) ) $req->set_param( 'per_page', (int) $args['per_page'] );
        $res = Semantic_Endpoints::entity_registry( $req );
        if ( is_wp_error( $res ) ) throw new \Exception( esc_html( $res->get_error_message() ) );
        return $res->data;
    }

    private static function tool_get_entity( array $args ) {
        $entity_id = (string) ( $args['entity_id'] ?? '' );
        if ( '' === $entity_id ) return self::tool_error( 'get_entity requires entity_id' );
        $row = Semantic_Entity_Registry::get_by_entity_id( $entity_id );
        if ( ! $row || ! Semantic_Entity_Registry::is_public( $row ) ) return self::tool_error( 'No entity found for id: ' . $entity_id );
        // public_only = true → linked_content is restricted to published, public pages.
        $detail = Semantic_Entity_Registry::get_detail( (int) $row['id'], true );
        return self::content_block( $detail ?: $row );
    }

    private static function tool_site_overview() {
        $req = new WP_REST_Request( 'GET', '/' . self::NS . '/organization' );
        $org = Semantic_Endpoints::organization( $req );
        $organization = is_wp_error( $org ) ? null : $org->data;

        // cluster_graph() is itself scoped to published, public content, so every
        // value below is derived only from currently-public pages (no draft leakage).
        $graph    = Semantic_Cohesion::cluster_graph();
        $clusters = array_map( fn( $c ) => [ 'label' => $c['label'], 'pages' => $c['size'] ], $graph['clusters'] ?? [] );

        // Dominant entities = the most frequently shared entities across the public
        // graph. Derived here (not from the stored cohesion snapshot) so this public
        // overview can never surface an entity that only appears in unpublished posts.
        $freq = [];
        foreach ( ( $graph['posts'] ?? [] ) as $p ) {
            foreach ( ( $p['entities'] ?? [] ) as $name ) {
                $freq[ $name ] = ( $freq[ $name ] ?? 0 ) + 1;
            }
        }
        arsort( $freq );
        $dominant = [];
        foreach ( array_slice( array_keys( $freq ), 0, 5 ) as $name ) {
            $dominant[] = [ 'name' => $name, 'frequency' => $freq[ $name ] ];
        }

        return [
            'organization'      => $organization,
            'topic_clusters'    => $clusters,
            'dominant_entities' => $dominant,
            'page_count'        => count( $graph['posts'] ?? [] ),
        ];
    }

    // -------------------------------------------------------------------------
    // Helpers
    // -------------------------------------------------------------------------

    private static function resolve_post( $id ) {
        // Public MCP surface — gate every resolution through is_publicly_servable()
        // (publish status, no password, not excluded, public type).
        if ( is_numeric( $id ) ) {
            $post = get_post( (int) $id );
            if ( $post && Semantic_Settings::is_publicly_servable( $post ) ) return $post;
        }
        $posts = get_posts( [
            'name'           => $id,
            'post_type'      => Semantic_Settings::post_types(),
            'post_status'    => 'publish',
            'posts_per_page' => 1,
        ] );
        $post = ! empty( $posts ) ? $posts[0] : null;
        return ( $post && Semantic_Settings::is_publicly_servable( $post ) ) ? $post : null;
    }

    /** Wrap a payload as a successful tools/call result (text content block). */
    private static function content_block( $payload ) {
        return [
            'content' => [ [
                'type' => 'text',
                'text' => wp_json_encode( $payload, JSON_UNESCAPED_SLASHES | JSON_UNESCAPED_UNICODE ),
            ] ],
        ];
    }

    /** Tool-level (not protocol-level) error. */
    private static function tool_error( $message ) {
        return [
            'content' => [ [ 'type' => 'text', 'text' => $message ] ],
            'isError' => true,
        ];
    }

    private static function ok( $id, $result ) {
        return [ 'jsonrpc' => '2.0', 'id' => $id, 'result' => $result ];
    }

    private static function error_obj( $id, $code, $message ) {
        return [ 'jsonrpc' => '2.0', 'id' => $id, 'error' => [ 'code' => $code, 'message' => $message ] ];
    }

    private static function notification_ack() {
        $r = new WP_REST_Response( null );
        $r->set_status( 202 );
        return self::with_headers( $r );
    }

    private static function json_response( $payload, $status ) {
        $r = new WP_REST_Response( $payload );
        $r->set_status( $status );
        $r->header( 'Content-Type', 'application/json; charset=utf-8' );
        return self::with_headers( $r );
    }

    /** Permissive CORS so the browser-based MCP Inspector can connect. */
    private static function with_headers( WP_REST_Response $r ) {
        $r->header( 'Access-Control-Allow-Origin', '*' );
        $r->header( 'Access-Control-Allow-Methods', 'POST, OPTIONS' );
        $r->header( 'Access-Control-Allow-Headers', 'Content-Type, Accept, Authorization, Mcp-Session-Id, MCP-Protocol-Version' );
        return $r;
    }
}