#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ListResourcesRequestSchema,
  ListToolsRequestSchema,
  ReadResourceRequestSchema,
} from '@modelcontextprotocol/sdk/types.js';
import fetch from 'node-fetch';
import { readFileSync, existsSync } from 'fs';
import { join, dirname } from 'path';
import { fileURLToPath } from 'url';

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

const REPO_ROOT = join(__dirname, '..');
const DEFAULT_MIRROR = join(REPO_ROOT, 'third-party', 'wormhole-ai-docs');
const ALLOWED_FETCH_PREFIX = 'https://wormhole.com/docs';
const LLMS_TXT_URL = `${ALLOWED_FETCH_PREFIX}/llms.txt`;

const CATEGORIES = [
  'basics',
  'ntt',
  'connect',
  'wtt',
  'settlement',
  'executor',
  'multigov',
  'queries',
  'transfer',
  'typescript-sdk',
  'solidity-sdk',
  'cctp',
  'reference',
];

function mirrorRoot() {
  return process.env.WORMHOLE_DOCS_MIRROR || DEFAULT_MIRROR;
}

function allowFetch() {
  return process.env.WORMHOLE_DOCS_FETCH === '1' || process.env.WORMHOLE_DOCS_FETCH === 'true';
}

function maxResourceBytes() {
  const n = parseInt(process.env.WORMHOLE_MAX_RESOURCE_BYTES || '5242880', 10);
  return Number.isFinite(n) && n > 0 ? n : 5242880;
}

/** @param {string} uri */
function uriToHttps(uri) {
  if (uri === 'wormhole://ai/llms.txt') return LLMS_TXT_URL;
  if (uri === 'wormhole://ai/site-index.json')
    return `${ALLOWED_FETCH_PREFIX}/ai/site-index.json`;
  if (uri === 'wormhole://ai/llms-full.jsonl')
    return `${ALLOWED_FETCH_PREFIX}/ai/llms-full.jsonl`;
  const m = uri.match(/^wormhole:\/\/ai\/categories\/([a-z0-9-]+\.md)$/);
  if (m) return `${ALLOWED_FETCH_PREFIX}/ai/categories/${m[1]}`;
  return null;
}

/** @param {string} uri */
function uriToRelativePath(uri) {
  if (uri === 'wormhole://ai/llms.txt') return 'llms.txt';
  if (uri === 'wormhole://ai/site-index.json') return 'site-index.json';
  if (uri === 'wormhole://ai/llms-full.jsonl') return 'llms-full.jsonl';
  const m = uri.match(/^wormhole:\/\/ai\/categories\/([a-z0-9-]+\.md)$/);
  if (m) return join('categories', m[1]);
  return null;
}

/** @param {string} url */
async function fetchAllowed(url) {
  if (!url || !url.startsWith(ALLOWED_FETCH_PREFIX)) {
    throw new Error(`Fetch URL not allowlisted: ${url}`);
  }
  const res = await fetch(url, {
    redirect: 'follow',
    headers: { 'user-agent': 'mcp-wormhole-docs/1.0' },
  });
  if (!res.ok) throw new Error(`HTTP ${res.status} for ${url}`);
  const len = res.headers.get('content-length');
  if (len && parseInt(len, 10) > 80 * 1024 * 1024) {
    throw new Error(`Refusing to fetch body larger than 80MB (${len} bytes)`);
  }
  return Buffer.from(await res.arrayBuffer());
}

/**
 * @param {string} uri
 */
async function readResourceContentAsync(uri) {
  const rel = uriToRelativePath(uri);
  if (!rel) throw new Error(`Unknown resource URI: ${uri}`);
  const root = mirrorRoot();
  const localPath = join(root, rel);
  let buf;
  if (existsSync(localPath)) buf = readFileSync(localPath);
  else if (allowFetch()) {
    const u = uriToHttps(uri);
    if (!u) throw new Error(`No HTTPS mapping for ${uri}`);
    buf = await fetchAllowed(u);
  } else {
    throw new Error(`Missing ${localPath}. Sync mirror or set WORMHOLE_DOCS_FETCH=1`);
  }
  return formatBuffer(uri, rel, buf);
}

/**
 * @param {string} uri
 * @param {string} rel
 * @param {Buffer} buf
 */
function formatBuffer(uri, rel, buf) {
  const isJsonl = rel === 'llms-full.jsonl';
  const max = maxResourceBytes();
  if (isJsonl && buf.length > max) {
    const head = buf.subarray(0, max).toString('utf8');
    return {
      mimeType: 'text/plain; charset=utf-8',
      text:
        `[Truncated: ${buf.length} bytes, showing first ${max} bytes. Set WORMHOLE_MAX_RESOURCE_BYTES or read ${join(mirrorRoot(), rel)} on disk for RAG.]\n\n` +
        head,
    };
  }

  if (rel.endsWith('.json')) {
    return { mimeType: 'application/json; charset=utf-8', text: buf.toString('utf8') };
  }
  if (rel.endsWith('.jsonl')) {
    return { mimeType: 'application/x-ndjson; charset=utf-8', text: buf.toString('utf8') };
  }
  return { mimeType: 'text/plain; charset=utf-8', text: buf.toString('utf8') };
}

function buildResourceList() {
  const resources = [
    {
      uri: 'wormhole://ai/llms.txt',
      name: 'Wormhole llms.txt',
      description: 'Tier 1: site map and links (llms.txt standard)',
      mimeType: 'text/plain',
    },
    {
      uri: 'wormhole://ai/site-index.json',
      name: 'Wormhole site-index.json',
      description: 'Tier 2: lightweight page index with previews',
      mimeType: 'application/json',
    },
    {
      uri: 'wormhole://ai/llms-full.jsonl',
      name: 'Wormhole llms-full.jsonl',
      description: 'Tier 4: full doc corpus (large; may truncate via MCP read)',
      mimeType: 'application/x-ndjson',
    },
  ];
  for (const c of CATEGORIES) {
    resources.push({
      uri: `wormhole://ai/categories/${c}.md`,
      name: `Wormhole category: ${c}`,
      description: `Tier 3: bundled markdown for ${c}`,
      mimeType: 'text/markdown',
    });
  }
  return resources;
}

async function loadSiteIndex() {
  const root = mirrorRoot();
  const p = join(root, 'site-index.json');
  if (existsSync(p)) {
    return JSON.parse(readFileSync(p, 'utf8'));
  }
  if (allowFetch()) {
    const buf = await fetchAllowed(`${ALLOWED_FETCH_PREFIX}/ai/site-index.json`);
    return JSON.parse(buf.toString('utf8'));
  }
  throw new Error(
    `site-index.json not found under ${root}. Run scripts/doc/sync-wormhole-ai-resources.sh or set WORMHOLE_DOCS_FETCH=1`
  );
}

/**
 * @param {string} query
 * @param {number} limit
 */
async function searchDocs(query, limit) {
  const q = (query || '').toLowerCase().trim();
  if (!q) return { results: [], note: 'Empty query' };

  const index = await loadSiteIndex();
  if (!Array.isArray(index)) {
    return { results: [], note: 'site-index.json is not an array' };
  }

  const scored = [];
  for (const page of index) {
    const title = (page.title || '').toLowerCase();
    const prev = (page.preview || '').toLowerCase();
    const slug = (page.slug || '').toLowerCase();
    const id = (page.id || '').toLowerCase();
    const cats = Array.isArray(page.categories) ? page.categories.join(' ').toLowerCase() : '';

    let score = 0;
    if (title.includes(q)) score += 10;
    if (slug.includes(q) || id.includes(q)) score += 8;
    if (prev.includes(q)) score += 5;
    if (cats.includes(q)) score += 3;
    for (const word of q.split(/\s+/).filter((w) => w.length > 2)) {
      if (title.includes(word)) score += 2;
      if (prev.includes(word)) score += 1;
    }

    if (score > 0) {
      scored.push({
        score,
        title: page.title,
        id: page.id,
        html_url: page.html_url,
        resolved_md_url: page.resolved_md_url,
        preview: page.preview
          ? page.preview.slice(0, 400) + (page.preview.length > 400 ? '…' : '')
          : undefined,
        categories: page.categories,
      });
    }
  }

  scored.sort((a, b) => b.score - a.score);
  return { results: scored.slice(0, limit) };
}

const server = new Server(
  { name: 'wormhole-docs', version: '1.0.0' },
  { capabilities: { resources: {}, tools: {} } }
);

server.setRequestHandler(ListResourcesRequestSchema, async () => ({
  resources: buildResourceList(),
}));

server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
  const uri = request.params.uri;
  const { text, mimeType } = await readResourceContentAsync(uri);
  return {
    contents: [{ uri, mimeType, text }],
  };
});

server.setRequestHandler(ListToolsRequestSchema, async () => ({
  tools: [
    {
      name: 'wormhole_doc_search',
      description:
        'Search Wormhole documentation site-index (titles, previews, categories). Use for tier-2 retrieval before loading full category markdown or llms-full.jsonl.',
      inputSchema: {
        type: 'object',
        properties: {
          query: { type: 'string', description: 'Search string (e.g. NTT, VAA, CCTP)' },
          limit: { type: 'number', description: 'Max results (default 10)', default: 10 },
        },
        required: ['query'],
      },
    },
  ],
}));

server.setRequestHandler(CallToolRequestSchema, async (request) => {
  if (request.params.name !== 'wormhole_doc_search') {
    throw new Error(`Unknown tool: ${request.params.name}`);
  }
  const args = request.params.arguments || {};
  const limit = Math.min(50, Math.max(1, parseInt(String(args.limit || 10), 10) || 10));
  const { results, note } = await searchDocs(String(args.query || ''), limit);
  return {
    content: [
      {
        type: 'text',
        text: JSON.stringify({ note, count: results.length, results }, null, 2),
      },
    ],
  };
});

const transport = new StdioServerTransport();
await server.connect(transport);