import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
import { createLogger } from '../../shared/log.ts'

const log = createLogger({ component: 'zai-mcp' })

/**
 * MCP bridge over the pinned @z_ai/mcp-server (vision).
 * Curated: only the mapped tool names below are exposed to the agent with
 * stable will-owned names; runtime `npx @latest` is forbidden (spec AC20).
 */

export const VISION_TOOL_MAP: Record<string, string> = {
  ui_to_artifact: 'zai_ui_to_artifact',
  extract_text_from_screenshot: 'zai_extract_text_from_screenshot',
  diagnose_error_screenshot: 'zai_diagnose_error_screenshot',
  understand_technical_diagram: 'zai_understand_technical_diagram',
  analyze_data_visualization: 'zai_analyze_data_visualization',
  ui_diff_check: 'zai_ui_diff_check',
  image_analysis: 'zai_image_analysis',
  video_analysis: 'zai_video_analysis',
}

/** Vision MCP documented limit: local/remote videos <= 8 MB, MP4/MOV/M4V. */
export const MCP_VIDEO_MAX_BYTES = 8 << 20

export interface McpBridge {
  listTools(): Promise<{ name: string; description?: string | undefined }[]>
  call(name: string, args: Record<string, unknown>): Promise<unknown>
  close(): Promise<void>
}

export async function startVisionMcp(opts: {
  command: string[]
  apiKey: string
  cwd: string
}): Promise<McpBridge> {
  const client = new Client({ name: 'will-agent', version: '0.1.0' }, { capabilities: {} })
  const transport = new StdioClientTransport({
    command: opts.command[0] ?? 'node',
    args: opts.command.slice(1),
    cwd: opts.cwd,
    env: {
      PATH: process.env.PATH ?? '',
      HOME: process.env.HOME ?? '',
      Z_AI_API_KEY: opts.apiKey,
      Z_AI_MODE: 'ZAI',
    },
  })
  await client.connect(transport)
  log.info('vision mcp connected', {})
  return {
    async listTools() {
      const res = await client.listTools()
      return res.tools.map((t) => ({
        name: t.name,
        ...(t.description !== undefined ? { description: t.description } : {}),
      }))
    },
    async call(name, args) {
      const res = await client.callTool({ name, arguments: args })
      return res.content
    },
    async close() {
      await client.close()
    },
  }
}