import { describe, it, expect, vi } from 'vitest';
import { sanitizeWireName, WireNameCodec, routeWithWireNames } from '../src/util/wire-names.js';
import type { JsonRpcRequest, JsonRpcResponse } from '../src/types.js';

describe('sanitizeWireName', () => {
  it('replaces slashes with underscores', () => {
    expect(sanitizeWireName('websearch/fetch_content')).toBe('websearch_fetch_content');
    expect(sanitizeWireName('all/websearch/fetch_content')).toBe('all_websearch_fetch_content');
  });

  it('keeps names that are already OpenAI-safe', () => {
    expect(sanitizeWireName('begin_session')).toBe('begin_session');
    expect(sanitizeWireName('my-grafana.tool')).toBe('my-grafana.tool');
  });

  it('replaces every character outside [A-Za-z0-9_.-]', () => {
    expect(sanitizeWireName('a b:c/d')).toBe('a_b_c_d');
  });
});

describe('WireNameCodec', () => {
  it('round-trips a namespaced tool name', () => {
    const codec = new WireNameCodec();
    const wire = codec.encodeName('websearch/fetch_content');
    expect(wire).toBe('websearch_fetch_content');
    expect(codec.decodeName(wire)).toBe('websearch/fetch_content');
  });

  it('is stable across repeated encodes', () => {
    const codec = new WireNameCodec();
    expect(codec.encodeName('searxng/web_url_read')).toBe('searxng_web_url_read');
    expect(codec.encodeName('searxng/web_url_read')).toBe('searxng_web_url_read');
  });

  it('passes unknown inbound names through unchanged', () => {
    const codec = new WireNameCodec();
    // Legacy client echoing the slash form, or a virtual tool never listed.
    expect(codec.decodeName('websearch/fetch_content')).toBe('websearch/fetch_content');
    expect(codec.decodeName('begin_session')).toBe('begin_session');
  });

  it('suffixes on collision, first registration wins the plain name', () => {
    const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
    try {
      const codec = new WireNameCodec();
      expect(codec.encodeName('foo_bar/baz')).toBe('foo_bar_baz');
      expect(codec.encodeName('foo/bar_baz')).toBe('foo_bar_baz_2');
      // Both decode back to their own presented names.
      expect(codec.decodeName('foo_bar_baz')).toBe('foo_bar/baz');
      expect(codec.decodeName('foo_bar_baz_2')).toBe('foo/bar_baz');
      // And stay stable.
      expect(codec.encodeName('foo/bar_baz')).toBe('foo_bar_baz_2');
      expect(warn).toHaveBeenCalledOnce();
    } finally {
      warn.mockRestore();
    }
  });

  it('encodes tools/list responses and leaves other fields intact', () => {
    const codec = new WireNameCodec();
    const response: JsonRpcResponse = {
      jsonrpc: '2.0',
      id: 1,
      result: {
        tools: [
          { name: 'websearch/fetch_content', description: 'fetch', inputSchema: { type: 'object' } },
          { name: 'begin_session', description: 'gate' },
        ],
      },
    };
    const encoded = codec.encodeToolsList(response);
    const tools = (encoded.result as { tools: Array<{ name: string; description?: string }> }).tools;
    expect(tools.map((t) => t.name)).toEqual(['websearch_fetch_content', 'begin_session']);
    expect(tools[0]?.description).toBe('fetch');
    // Original response object is not mutated.
    const originalTools = (response.result as { tools: Array<{ name: string }> }).tools;
    expect(originalTools[0]?.name).toBe('websearch/fetch_content');
  });

  it('returns error and non-list responses unchanged', () => {
    const codec = new WireNameCodec();
    const err: JsonRpcResponse = { jsonrpc: '2.0', id: 1, error: { code: -32603, message: 'boom' } };
    expect(codec.encodeToolsList(err)).toBe(err);
    const other: JsonRpcResponse = { jsonrpc: '2.0', id: 1, result: { content: [] } };
    expect(codec.encodeToolsList(other)).toBe(other);
  });

  it('decodes tools/call requests for known wire names only', () => {
    const codec = new WireNameCodec();
    codec.encodeName('websearch/fetch_content');

    const known: JsonRpcRequest = {
      jsonrpc: '2.0',
      id: 2,
      method: 'tools/call',
      params: { name: 'websearch_fetch_content', arguments: { url: 'https://x' } },
    };
    const decoded = codec.decodeToolCall(known);
    expect(decoded.params?.['name']).toBe('websearch/fetch_content');
    expect(decoded.params?.['arguments']).toEqual({ url: 'https://x' });
    // Original request object is not mutated.
    expect(known.params?.['name']).toBe('websearch_fetch_content');

    const unknown: JsonRpcRequest = {
      jsonrpc: '2.0',
      id: 3,
      method: 'tools/call',
      params: { name: 'not_listed', arguments: {} },
    };
    expect(codec.decodeToolCall(unknown)).toBe(unknown);
  });
});

describe('routeWithWireNames', () => {
  const listResponse: JsonRpcResponse = {
    jsonrpc: '2.0',
    id: 1,
    result: { tools: [{ name: 'websearch/fetch_content' }, { name: 'searxng/web_url_read' }] },
  };

  it('serves wire-safe names on tools/list and maps tools/call back', async () => {
    const codec = new WireNameCodec();
    const seen: JsonRpcRequest[] = [];
    const route = async (req: JsonRpcRequest): Promise<JsonRpcResponse> => {
      seen.push(req);
      if (req.method === 'tools/list') return listResponse;
      return { jsonrpc: '2.0', id: req.id, result: { content: [{ type: 'text', text: 'ok' }] } };
    };

    const listed = await routeWithWireNames(codec, route, { jsonrpc: '2.0', id: 1, method: 'tools/list' });
    const names = (listed.result as { tools: Array<{ name: string }> }).tools.map((t) => t.name);
    expect(names).toEqual(['websearch_fetch_content', 'searxng_web_url_read']);

    // The exact scenario from the librechat incident: the model echoes the
    // wire name; the router must receive the canonical name.
    await routeWithWireNames(codec, route, {
      jsonrpc: '2.0',
      id: 2,
      method: 'tools/call',
      params: { name: 'websearch_fetch_content', arguments: { url: 'https://x' } },
    });
    expect(seen[1]?.params?.['name']).toBe('websearch/fetch_content');
  });

  it('keeps legacy slash-name calls working', async () => {
    const codec = new WireNameCodec();
    const seen: JsonRpcRequest[] = [];
    const route = async (req: JsonRpcRequest): Promise<JsonRpcResponse> => {
      seen.push(req);
      return { jsonrpc: '2.0', id: req.id, result: {} };
    };
    await routeWithWireNames(codec, route, {
      jsonrpc: '2.0',
      id: 1,
      method: 'tools/call',
      params: { name: 'websearch/fetch_content', arguments: {} },
    });
    expect(seen[0]?.params?.['name']).toBe('websearch/fetch_content');
  });

  it('does not touch other methods', async () => {
    const codec = new WireNameCodec();
    const init: JsonRpcRequest = { jsonrpc: '2.0', id: 1, method: 'initialize', params: {} };
    const response: JsonRpcResponse = { jsonrpc: '2.0', id: 1, result: { protocolVersion: '2024-11-05' } };
    const out = await routeWithWireNames(codec, async () => response, init);
    expect(out).toBe(response);
  });
});