import { describe, it, expect, afterAll } from "bun:test";
import type { Subprocess } from "bun";

interface JsonRpcRequest {
  jsonrpc: "2.0";
  id?: number;
  method: string;
  params?: Record<string, unknown>;
}

interface JsonRpcResponse {
  jsonrpc: "2.0";
  id: number;
  result?: unknown;
  error?: { code: number; message: string; data?: unknown };
}

class McpTestClient {
  private proc: Subprocess<"pipe", "pipe", "pipe">;
  private buffer = "";
  private initialized = false;

  constructor() {
    this.proc = Bun.spawn(["bun", "run", "src/index.ts"], {
      stdin: "pipe",
      stdout: "pipe",
      stderr: "pipe",
      env: { ...process.env, SAGE_X3_MODE: "sandbox" },
    });
  }

  private async send(message: JsonRpcRequest): Promise<void> {
    const data = JSON.stringify(message) + "\n";
    this.proc.stdin.write(data);
    await this.proc.stdin.flush();
  }

  // Reads chunks from stdout until a complete JSON-RPC line is found.
  // Uses racing with timeout to avoid blocking forever on a hung process.
  private async readResponse(timeoutMs = 10000): Promise<JsonRpcResponse> {
    const deadline = Date.now() + timeoutMs;
    const reader = this.proc.stdout.getReader();

    try {
      while (Date.now() < deadline) {
        const newlineIdx = this.buffer.indexOf("\n");
        if (newlineIdx !== -1) {
          const line = this.buffer.slice(0, newlineIdx).trim();
          this.buffer = this.buffer.slice(newlineIdx + 1);
          if (line.length > 0) {
            try {
              return JSON.parse(line) as JsonRpcResponse;
            } catch {
              continue;
            }
          }
          continue;
        }

        const result = await Promise.race([
          reader.read(),
          new Promise<{ done: true; value: undefined }>((resolve) =>
            setTimeout(() => resolve({ done: true, value: undefined }), Math.max(100, deadline - Date.now())),
          ),
        ]);

        if (result.done && !result.value) continue;
        if (result.value) {
          this.buffer += new TextDecoder().decode(result.value);
        }
        if (result.done) break;
      }
    } finally {
      reader.releaseLock();
    }

    throw new Error(`Timeout waiting for JSON-RPC response after ${timeoutMs}ms. Buffer: ${this.buffer.slice(0, 200)}`);
  }

  async initialize(): Promise<JsonRpcResponse> {
    await this.send({
      jsonrpc: "2.0",
      id: 1,
      method: "initialize",
      params: {
        protocolVersion: "2024-11-05",
        capabilities: {},
        clientInfo: { name: "test", version: "1.0.0" },
      },
    });
    const response = await this.readResponse();

    // notifications/initialized has no id → no response expected
    await this.send({ jsonrpc: "2.0", method: "notifications/initialized" });
    await Bun.sleep(100);

    this.initialized = true;
    return response;
  }

  async request(id: number, method: string, params?: Record<string, unknown>): Promise<JsonRpcResponse> {
    if (!this.initialized) {
      throw new Error("Must call initialize() before sending requests");
    }
    await this.send({ jsonrpc: "2.0", id, method, params });
    return this.readResponse();
  }

  async close(): Promise<void> {
    try { this.proc.stdin.end(); } catch { /* already closed */ }
    this.proc.kill();
    await Promise.race([this.proc.exited, Bun.sleep(3000)]);
  }
}

describe("MCP Server Protocol", () => {
  let client: McpTestClient;

  afterAll(async () => {
    if (client) await client.close();
  });

  it("starts and responds to initialize handshake", async () => {
    client = new McpTestClient();
    const response = await client.initialize();

    expect(response.jsonrpc).toBe("2.0");
    expect(response.id).toBe(1);
    expect(response.error).toBeUndefined();
    expect(response.result).toBeDefined();

    const result = response.result as {
      protocolVersion: string;
      capabilities: Record<string, unknown>;
      serverInfo: { name: string; version: string };
    };

    expect(result.protocolVersion).toBeDefined();
    expect(result.serverInfo).toBeDefined();
    expect(result.serverInfo.name).toBe("sage-x3-graphql");
    expect(result.serverInfo.version).toBe("0.1.0");
    expect(result.capabilities).toBeDefined();
  }, 15000);

  it("returns 5 tools from tools/list", async () => {
    const response = await client.request(2, "tools/list");

    expect(response.jsonrpc).toBe("2.0");
    expect(response.id).toBe(2);
    expect(response.error).toBeUndefined();

    const result = response.result as { tools: Array<{ name: string; description: string }> };
    expect(result.tools).toBeDefined();
    expect(Array.isArray(result.tools)).toBe(true);
    expect(result.tools.length).toBe(5);

    const toolNames = result.tools.map((t) => t.name).sort();
    expect(toolNames).toEqual([
      "aggregate_entities",
      "execute_graphql",
      "introspect_schema",
      "query_entities",
      "read_entity",
    ]);

    for (const tool of result.tools) {
      expect(tool.description).toBeTruthy();
    }
  }, 10000);

  it("returns 5 resources from resources/list", async () => {
    const response = await client.request(3, "resources/list");

    expect(response.jsonrpc).toBe("2.0");
    expect(response.id).toBe(3);
    expect(response.error).toBeUndefined();

    const result = response.result as {
      resources: Array<{ name: string; uri: string; description?: string }>;
    };
    expect(result.resources).toBeDefined();
    expect(Array.isArray(result.resources)).toBe(true);
    expect(result.resources.length).toBe(5);

    for (const resource of result.resources) {
      expect(resource.uri).toMatch(/^sage-x3:\/\//);
      expect(resource.name).toBeTruthy();
    }
  }, 10000);

  it("returns 4 prompts from prompts/list", async () => {
    const response = await client.request(4, "prompts/list");

    expect(response.jsonrpc).toBe("2.0");
    expect(response.id).toBe(4);
    expect(response.error).toBeUndefined();

    const result = response.result as {
      prompts: Array<{ name: string; description?: string }>;
    };
    expect(result.prompts).toBeDefined();
    expect(Array.isArray(result.prompts)).toBe(true);
    expect(result.prompts.length).toBe(4);

    const promptNames = result.prompts.map((p) => p.name).sort();
    expect(promptNames).toEqual([
      "analyze-data",
      "explore-data",
      "lookup-entity",
      "search-entities",
    ]);
  }, 10000);

  it("handles unknown method with error response", async () => {
    const response = await client.request(5, "nonexistent/method");

    expect(response.jsonrpc).toBe("2.0");
    expect(response.id).toBe(5);
    expect(response.error).toBeDefined();
    expect(response.error!.code).toBeDefined();
    expect(response.error!.message).toBeTruthy();
  }, 10000);

  it("tools have valid input schemas", async () => {
    const response = await client.request(6, "tools/list");
    const result = response.result as {
      tools: Array<{ name: string; inputSchema: Record<string, unknown> }>;
    };

    for (const tool of result.tools) {
      expect(tool.inputSchema).toBeDefined();
      expect(tool.inputSchema.type).toBe("object");
    }
  }, 10000);
});