import { afterEach, beforeEach, describe, expect, test } from 'bun:test'
import type { AddressInfo } from 'node:net'

import type { Client } from '@modelcontextprotocol/sdk/client/index.js'
import type { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
import { WebSocketServer, type WebSocket } from 'ws'

import {
  ALL_TOOLS,
  FigmaAPI,
  SceneGraph,
  computeAllLayouts,
  executeRpcCommand
} from '@open-pencil/core'

import { expectDefined, getNodeOrThrow } from '#tests/helpers/assert'

function createMockApp() {
  const graph = new SceneGraph()
  const wss = new WebSocketServer({ port: 0, host: '127.0.0.1' })
  let clientWs: WebSocket | null = null
  const requests: Array<{
    command: string
    args?: { name?: string; document_id?: string; page_id?: string; args?: Record<string, unknown> }
  }> = []

  wss.on('connection', (ws) => {
    clientWs = ws
    ws.send(JSON.stringify({ type: 'register', token: 'mock-token' }))

    ws.on('message', async (raw) => {
      const msg = JSON.parse(String(raw)) as {
        type: string
        id: string
        command: string
        args?: { name?: string; args?: Record<string, unknown> }
      }
      if (msg.type !== 'request') return
      requests.push({ command: msg.command, args: msg.args })

      try {
        let result: unknown
        if (msg.command === 'tool' && msg.args?.name) {
          const toolName = msg.args.name
          const def = ALL_TOOLS.find((t) => t.name === toolName)
          if (!def) throw new Error(`Unknown tool: ${toolName}`)
          const api = new FigmaAPI(graph)
          api.currentPage = api.wrapNode(graph.getPages()[0].id)
          result = await def.execute(api, msg.args.args ?? {})
          if (def.mutates) computeAllLayouts(graph)
        } else if (msg.command === 'save_file') {
          result = { ok: true }
        } else if (msg.command === 'list_documents') {
          result = {
            documents: [
              {
                id: 'doc-1',
                name: 'Mock document',
                active: true,
                current_page_id: graph.getPages()[0].id,
                current_page_name: graph.getPages()[0].name,
                pages: graph.getPages().map((page) => ({ id: page.id, name: page.name }))
              }
            ]
          }
        } else {
          result = executeRpcCommand(graph, msg.command, msg.args ?? {})
        }

        ws.send(JSON.stringify({ type: 'response', id: msg.id, ok: true, result }))
      } catch (e) {
        ws.send(
          JSON.stringify({
            type: 'response',
            id: msg.id,
            ok: false,
            error: e instanceof Error ? e.message : String(e)
          })
        )
      }
    })
  })

  const port = new Promise<number>((resolve) => {
    wss.on('listening', () => resolve((wss.address() as AddressInfo).port))
  })

  return {
    graph,
    requests,
    wss,
    port,
    close: () => {
      clientWs?.close()
      wss.close()
    }
  }
}

async function createStdioClient(wsPort: number) {
  const { Client } = await import('@modelcontextprotocol/sdk/client/index.js')
  const { StdioClientTransport } = await import('@modelcontextprotocol/sdk/client/stdio.js')
  const transport = new StdioClientTransport({
    command: 'bun',
    args: ['packages/mcp/src/stdio.ts'],
    env: {
      ...process.env,
      WS_PORT: String(wsPort),
      PATH: process.env.PATH ?? ''
    },
    stderr: 'pipe'
  })

  const client = new Client({ name: 'test-stdio-client', version: '0.0.0' })

  await new Promise<void>((resolve) => {
    const stderr = transport.stderr
    if (stderr && 'on' in stderr) {
      ;(stderr as NodeJS.ReadableStream).on('data', (chunk: Buffer) => {
        if (chunk.toString().includes('Connected to OpenPencil app')) {
          resolve()
        }
      })
    }
    void client.connect(transport).then(() => {
      setTimeout(resolve, 1000)
      return undefined
    })
  })

  return { client, transport }
}

function textContent(content: unknown): string {
  const items = content as { type: string; text: string }[]
  return expectDefined(
    items.find((c) => c.type === 'text'),
    'text content'
  ).text
}

describe('MCP stdio transport', () => {
  let app: ReturnType<typeof createMockApp>
  let client: Client
  let transport: StdioClientTransport

  beforeEach(async () => {
    app = createMockApp()
    const wsPort = await app.port
    const ctx = await createStdioClient(wsPort)
    client = ctx.client
    transport = ctx.transport
  })

  afterEach(async () => {
    await client.close()
    app.close()
  })

  test('lists tools over stdio', async () => {
    const { tools } = await client.listTools()
    const names = tools.map((t) => t.name)
    expect(names).toContain('create_shape')
    expect(names).toContain('get_page_tree')
    expect(names).toContain('save_file')
    expect(names).toContain('list_documents')
    expect(names).toContain('get_codegen_prompt')
    const createShape = expectDefined(
      tools.find((tool) => tool.name === 'create_shape'),
      'create_shape tool'
    )
    expect(JSON.stringify(createShape.inputSchema)).toContain('document_id')
    expect(JSON.stringify(createShape.inputSchema)).toContain('page_id')
    expect(tools.length).toBeGreaterThan(30)
  })

  test('create_shape via stdio creates a node', async () => {
    const result = await client.callTool({
      name: 'create_shape',
      arguments: { type: 'FRAME', x: 10, y: 20, width: 200, height: 100, name: 'StdioFrame' }
    })
    expect(result.isError).not.toBe(true)
    const data = JSON.parse(textContent(result.content)) as {
      id: string
      name: string
      type: string
    }
    expect(data.type).toBe('FRAME')
    expect(data.name).toBe('StdioFrame')

    expect(getNodeOrThrow(app.graph, data.id).width).toBe(200)
  })

  test('tool target fields are sent in the app RPC envelope', async () => {
    const result = await client.callTool({
      name: 'create_shape',
      arguments: {
        document_id: 'doc-1',
        page_id: 'page-1',
        type: 'FRAME',
        x: 10,
        y: 20,
        width: 200,
        height: 100,
        name: 'TargetedFrame'
      }
    })
    expect(result.isError).not.toBe(true)
    const request = expectDefined(
      app.requests.find((item) => item.command === 'tool' && item.args?.name === 'create_shape'),
      'tool request'
    )
    expect(request.args?.document_id).toBe('doc-1')
    expect(request.args?.page_id).toBe('page-1')
    expect(request.args?.args?.document_id).toBeUndefined()
    expect(request.args?.args?.page_id).toBeUndefined()
  })

  test('list_documents via stdio returns open documents', async () => {
    const result = await client.callTool({ name: 'list_documents', arguments: {} })
    expect(result.isError).not.toBe(true)
    const data = JSON.parse(textContent(result.content)) as {
      documents: Array<{ id: string; current_page_id: string }>
    }
    expect(data.documents[0].id).toBe('doc-1')
    expect(data.documents[0].current_page_id).toBe(app.graph.getPages()[0].id)
  })

  test('save_file via stdio succeeds', async () => {
    const result = await client.callTool({ name: 'save_file', arguments: {} })
    expect(result.isError).not.toBe(true)
    const data = JSON.parse(textContent(result.content)) as { saved: boolean }
    expect(data.saved).toBe(true)
  })

  test('get_codegen_prompt via stdio returns prompt', async () => {
    const result = await client.callTool({ name: 'get_codegen_prompt', arguments: {} })
    expect(result.isError).not.toBe(true)
    const data = JSON.parse(textContent(result.content)) as { prompt: string }
    expect(data.prompt.length).toBeGreaterThan(100)
  })

  test('delete_node via stdio removes a node', async () => {
    const create = await client.callTool({
      name: 'create_shape',
      arguments: { type: 'RECTANGLE', x: 0, y: 0, width: 50, height: 50 }
    })
    const { id } = JSON.parse(textContent(create.content)) as { id: string }

    expect(app.graph.getNode(id)).toBeDefined()

    await client.callTool({ name: 'delete_node', arguments: { id } })

    expect(app.graph.getNode(id)).toBeUndefined()
  })

  test('stderr does not contain JSON-RPC', async () => {
    const stderrChunks: string[] = []
    const stderr = transport.stderr
    if (stderr && 'on' in stderr) {
      ;(stderr as NodeJS.ReadableStream).on('data', (chunk: Buffer) => {
        stderrChunks.push(chunk.toString())
      })
    }

    await client.callTool({
      name: 'create_shape',
      arguments: { type: 'FRAME', x: 0, y: 0, width: 100, height: 100 }
    })

    const allStderr = stderrChunks.join('')
    expect(allStderr).not.toContain('"jsonrpc"')
    expect(allStderr).not.toContain('"method"')
  })
})