import assert from 'node:assert/strict'
import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises'
import { createServer, type Server } from 'node:http'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { test } from 'node:test'
import { startService } from '../service/main.ts'
import { createServiceClient } from './service-client.ts'
import type { Tool } from './tools-shape.ts'
import { ZaiRestClient } from './zai/client.ts'
import { startVisionMcp, VISION_TOOL_MAP } from './zai/mcp.ts'
import { createZaiTools } from './zai/tools.ts'

interface FakeZai {
  server: Server
  url: string
  calls: { path: string; body?: string }[]
  failNext: number
}

async function startFakeZai(): Promise<FakeZai> {
  const calls: { path: string; body?: string }[] = []
  const state = { failNext: 0 }
  const server = createServer((req, res) => {
    const url = new URL(req.url ?? '/', 'http://x')
    let data = ''
    req.on('data', (c: Buffer) => {
      data += c.toString('utf8')
    })
    req.on('end', () => {
      calls.push({ path: url.pathname, body: data.slice(0, 200) })
      if (state.failNext > 0) {
        state.failNext--
        res.writeHead(500)
        res.end('{}')
        return
      }
      const json = (payload: unknown) => {
        res.writeHead(200, { 'content-type': 'application/json' })
        res.end(JSON.stringify(payload))
      }
      if (url.pathname === '/paas/v4/chat/completions') {
        if (data.includes('"model":"bad"')) {
          res.writeHead(400, { 'content-type': 'application/json' })
          res.end(JSON.stringify({ error: { message: 'bad model' } }))
          return
        }
        return json({
          id: 'cmpl-1',
          choices: [{ message: { content: 'a cat sleeping on a keyboard' } }],
          usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
        })
      }
      if (url.pathname === '/paas/v4/audio/transcriptions') {
        return json({ text: 'hello from audio' })
      }
      if (url.pathname === '/paas/v4/layout_parsing') {
        return json({ md_results: '# parsed doc', data_info: { num_pages: 2 } })
      }
      if (url.pathname === '/paas/v4/images/generations') {
        return json({
          id: 'img-1',
          data: [{ b64_json: Buffer.from('png-bytes').toString('base64') }],
        })
      }
      if (url.pathname === '/paas/v4/videos/generations') {
        return json({ id: 'vid-1' })
      }
      if (url.pathname.startsWith('/paas/v4/videos/retrieve')) {
        return json({ task_status: 'SUCCESS', video_url: `${base}/files/out.mp4` })
      }
      if (url.pathname === '/files/out.mp4') {
        res.writeHead(200, { 'content-type': 'video/mp4' })
        res.end(Buffer.from('mp4-bytes'))
        return
      }
      res.writeHead(404)
      res.end('{}')
    })
  })
  await new Promise<void>((r) => server.listen(0, '127.0.0.1', () => r()))
  const base = `http://127.0.0.1:${(server.address() as { port: number }).port}`
  return {
    server,
    url: base,
    calls,
    get failNext() {
      return state.failNext
    },
    set failNext(v: number) {
      state.failNext = v
    },
  }
}

async function toolByName(tools: readonly unknown[], name: string): Promise<Tool> {
  const found = tools.find((t) => (t as { name?: string }).name === name) as Tool | undefined
  if (!found) throw new Error(`tool ${name} missing`)
  return found
}

async function runTool(tool: Tool, args: Record<string, unknown>): Promise<unknown> {
  const res = await tool.execute('call-1', args, new AbortController().signal, () => {}, {
    cwd: process.cwd(),
  })
  const text =
    (res.content as { type: string; text?: string }[]).find((c) => c.type === 'text')?.text ?? ''
  return JSON.parse(text)
}

async function withService<T>(
  fn: (client: ReturnType<typeof createServiceClient>) => Promise<T>,
): Promise<T> {
  const dir = await mkdtemp(join(tmpdir(), 'zai-'))
  const webDist = join(dir, 'web-dist')
  await mkdir(webDist, { recursive: true })
  const service = await startService({
    dataDir: dir,
    publicAddr: '127.0.0.1:0',
    internalAddr: '127.0.0.1:0',
    webDist,
    docsDist: webDist,
    mediaMaxBytes: 100 << 20,
    storageWarnPct: 20,
    storageStopPct: 10,
    workerUrl: join(import.meta.dirname, '../service/worker.ts'),
  })
  try {
    const addr = service.internalServer.address()
    if (addr === null || typeof addr !== 'object') throw new Error('no addr')
    return await fn(createServiceClient(`http://127.0.0.1:${addr.port}`))
  } finally {
    await service.close()
  }
}

test('rest client retries retryable failures', async () => {
  const fake = await startFakeZai()
  const client = new ZaiRestClient('KEY', fake.url, { retries: 2, timeoutMs: 5000 })
  fake.failNext = 1
  const res = await client.chatCompletion({
    model: 'glm-5.3-flash',
    messages: [{ role: 'user', content: 'hi' }],
  })
  assert.equal(res.content, 'a cat sleeping on a keyboard')
  assert.equal(fake.calls.filter((c) => c.path === '/paas/v4/chat/completions').length, 2)
  fake.server.close()
})

test('rest client surfaces 4xx without retry', async () => {
  const fake = await startFakeZai()
  const client = new ZaiRestClient('KEY', fake.url, { retries: 3, timeoutMs: 5000 })
  await assert.rejects(
    client.chatCompletion({ model: 'bad', messages: [{ role: 'user', content: 'hi' }] }),
    (err: unknown) => {
      assert.ok(err instanceof Error)
      assert.match(err.message, /400/)
      return true
    },
  )
  assert.equal(fake.calls.filter((c) => c.path === '/paas/v4/chat/completions').length, 1)
  fake.server.close()
})

test('zai_analyze_media answers, caches by digest, second call is a hit', async () => {
  const fake = await startFakeZai()
  await withService(async (service) => {
    const meta = await service.putMedia(new Uint8Array(Buffer.from('imgdata')), {
      type: 'image/png',
    })
    const client = new ZaiRestClient('KEY', fake.url, { timeoutMs: 5000 })
    const tools = createZaiTools({ client, service, analysisModel: 'glm-5.3-flash' })
    const tool = await toolByName(tools, 'zai_analyze_media')
    const first = (await runTool(tool, { digest: meta.digest, question: 'what is it?' })) as {
      ok: boolean
      fromCache: boolean
      answer: string
    }
    assert.equal(first.ok, true)
    assert.equal(first.fromCache, false)
    assert.match(first.answer, /cat/)
    const second = (await runTool(tool, { digest: meta.digest, question: 'again?' })) as {
      fromCache: boolean
    }
    assert.equal(second.fromCache, true)
    assert.equal(fake.calls.filter((c) => c.path === '/paas/v4/chat/completions').length, 1)
  })
  fake.server.close()
})

test('zai_transcribe and zai_parse_document return explicit results', async () => {
  const fake = await startFakeZai()
  await withService(async (service) => {
    const audio = await service.putMedia(new Uint8Array(Buffer.from('wav')), { type: 'audio/wav' })
    const doc = await service.putMedia(new Uint8Array(Buffer.from('%PDF-1.4')), {
      type: 'application/pdf',
    })
    const client = new ZaiRestClient('KEY', fake.url, { timeoutMs: 5000 })
    const tools = createZaiTools({ client, service, analysisModel: 'glm-5.3-flash' })
    const tr = await runTool(await toolByName(tools, 'zai_transcribe'), { digest: audio.digest })
    assert.deepEqual(tr, { ok: true, fromCache: false, text: 'hello from audio' })
    const pd = await runTool(await toolByName(tools, 'zai_parse_document'), { digest: doc.digest })
    assert.deepEqual(pd, { ok: true, fromCache: false, markdown: '# parsed doc', pages: 2 })
  })
  fake.server.close()
})

test('generation tools store bytes and return digests; limits fail explicitly', async () => {
  const fake = await startFakeZai()
  await withService(async (service) => {
    const client = new ZaiRestClient('KEY', fake.url, { timeoutMs: 5000 })
    const tools = createZaiTools({ client, service, analysisModel: 'glm-5.3-flash' })
    const img = (await runTool(await toolByName(tools, 'zai_generate_image'), {
      prompt: 'a red cube',
    })) as { ok: boolean; digest: string }
    assert.equal(img.ok, true)
    const imgBytes = await service.getMedia(img.digest)
    assert.deepEqual(Buffer.from(imgBytes).toString(), 'png-bytes')

    const vid = (await runTool(await toolByName(tools, 'zai_generate_video'), {
      prompt: 'waves',
    })) as { ok: boolean; digest: string }
    assert.equal(vid.ok, true)
    const vidBytes = await service.getMedia(vid.digest)
    assert.deepEqual(Buffer.from(vidBytes).toString(), 'mp4-bytes')

    const big = await service.putMedia(new Uint8Array(11 << 20), { type: 'image/png' })
    const tooBig = (await runTool(await toolByName(tools, 'zai_analyze_media'), {
      digest: big.digest,
      question: 'x',
    })) as { ok: boolean; error: string }
    assert.equal(tooBig.ok, false)
    assert.match(tooBig.error, /10 MB/)
  })
  fake.server.close()
})

test('curated tool set carries stable names and one-line descriptions', async () => {
  const fake = await startFakeZai()
  await withService(async (service) => {
    const client = new ZaiRestClient('KEY', fake.url, { timeoutMs: 5000 })
    const tools = createZaiTools({ client, service, analysisModel: 'glm-5.3-flash' })
    assert.deepEqual(tools.map((t) => t.name).sort(), [
      'zai_analyze_media',
      'zai_generate_image',
      'zai_generate_video',
      'zai_parse_document',
      'zai_transcribe',
    ])
    for (const t of tools) {
      assert.ok(t.description && t.description.length > 20, `${t.name} needs a real description`)
    }
  })
  fake.server.close()
})

test('mcp bridge lists and calls tools over stdio', async () => {
  // Fake MCP stdio server implemented with the official SDK server class.
  // Kept inside the repo tree so the script can resolve node_modules.
  const script = join(import.meta.dirname, '__fake-mcp-server.mjs')
  await writeFile(
    script,
    `import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'
import { z } from 'zod'
const server = new McpServer({ name: 'fake-zai', version: '0.0.1' })
server.tool('image_analysis', { image_path: z.string(), prompt: z.string() }, async ({ image_path, prompt }) => ({
  content: [{ type: 'text', text: 'analysis of ' + image_path + ': ' + prompt }],
}))
await server.connect(new StdioServerTransport())
`,
    'utf8',
  )
  const bridge = await startVisionMcp({
    command: [process.execPath, script],
    apiKey: 'KEY',
    cwd: process.cwd(),
  })
  try {
    const tools = await bridge.listTools()
    assert.equal(tools[0]?.name, 'image_analysis')
    const result = await bridge.call('image_analysis', { image_path: 'a.png', prompt: 'what' })
    assert.match(JSON.stringify(result), /analysis of a\.png/)
    assert.ok(VISION_TOOL_MAP.image_analysis === 'zai_image_analysis')
  } finally {
    await bridge.close()
    await rm(script, { force: true })
  }
})