import test from 'node:test'
import assert from 'node:assert/strict'
import { EventEmitter } from 'node:events'
import { PassThrough } from 'node:stream'
import { mkdtempSync, writeFileSync, readFileSync, mkdirSync, existsSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { createCaptionHost, validateCaptionHostRequest, captionModelFiles } from '../scripts/caption-host.mjs'
import { validateCaptionRequest, captionStylePrompt, CAPTION_STYLES, DELTA_LOCKED_LOOK } from '../shared/studio2/caption.mjs'

test('caption styles and delta lock text', () => {
  assert.deepEqual(CAPTION_STYLES, ['descriptive', 'klein_prompt', 'delta', 'tags'])
  assert.match(captionStylePrompt('delta'), new RegExp(DELTA_LOCKED_LOOK))
  assert.match(captionStylePrompt('klein_prompt'), /Flux Klein/)
  assert.match(captionStylePrompt('tags'), /comma-separated/)
})

test('validateCaptionRequest requires folder + still', () => {
  assert.throws(() => validateCaptionRequest({}), /folder/)
  assert.throws(() => validateCaptionRequest({ folderId: 'f1' }), /still/)
  assert.throws(() => validateCaptionRequest({ folderId: 'f1', stillId: 's1', captionStyle: 'nope' }), /style/)
  const ok = validateCaptionRequest({ folderId: 'f1', stillId: 's1', captionStyle: 'descriptive' })
  assert.equal(ok.mode, 'caption')
  assert.equal(ok.captionStyle, 'descriptive')
  assert.equal(ok.lockFace, false)
  assert.equal(ok.lockOutfit, false)
})

test('host validates job body', () => {
  assert.throws(() => validateCaptionHostRequest({ id: 'bad', style: 'descriptive' }))
  assert.deepEqual(
    validateCaptionHostRequest({ id: 'caption-job-123456', style: 'tags' }),
    { id: 'caption-job-123456', style: 'tags', imagePath: '' }
  )
})

function fixture() {
  const root = mkdtempSync(join(tmpdir(), 'aigen-caption-'))
  const models = join(root, 'models')
  mkdirSync(models)
  writeFileSync(join(models, captionModelFiles.model), 'model')
  writeFileSync(join(models, captionModelFiles.mmproj), 'mmproj')
  const llama = join(root, 'llama-server.exe')
  writeFileSync(llama, '')
  const dataDir = join(root, 'jobs')
  const child = new EventEmitter()
  child.pid = 4242
  child.stdout = new PassThrough()
  child.stderr = new PassThrough()
  child.kill = () => { setImmediate(() => child.emit('close', 0)); return true }
  child.exitCode = null
  const calls = []
  let fetchCount = 0
  const fetchImpl = async (url, init) => {
    fetchCount++
    if (String(url).endsWith('/health')) return { ok: true }
    if (String(url).includes('/v1/chat/completions')) {
      assert.match(JSON.stringify(JSON.parse(init.body)), /Describe this image|Flux Klein|locked baseline|comma-separated/)
      return {
        ok: true,
        json: async () => ({ choices: [{ message: { content: 'a warm caption' } }] })
      }
    }
    throw new Error(`unexpected fetch ${url}`)
  }
  const host = createCaptionHost({
    modelsDir: models,
    llamaServer: llama,
    dataDir,
    prepare: async () => {},
    leaseValid: () => true,
    spawnProcess: (...args) => { calls.push(args); setImmediate(() => {}); return child },
    fetchImpl
  })
  return { host, child, calls, root, dataDir, fetchCount: () => fetchCount }
}

test('host loads llama-server, captions once, and unloads', async () => {
  const f = fixture()
  const jobDir = join(f.dataDir, 'caption-job-123456')
  mkdirSync(jobDir, { recursive: true })
  writeFileSync(join(jobDir, 'input.png'), Buffer.alloc(64, 1))
  const started = f.host.start({ id: 'caption-job-123456', style: 'descriptive' }, 'lease')
  // Let the host reach waitForServer / chat completion.
  await new Promise(r => setTimeout(r, 50))
  f.child.emit('spawn')
  const result = await started
  assert.equal(result.status, 'complete')
  assert.equal(result.text, 'a warm caption')
  assert.equal(f.calls.length, 1)
  const [, args] = f.calls[0]
  assert.ok(args.includes('--mmproj'))
  assert.ok(args.includes('-m'))
  assert.equal(f.host.read('caption-job-123456').text, 'a warm caption')
  // busy may briefly remain true during the post-unload hold window
  assert.ok(typeof f.host.busy() === 'boolean')
})

test('host rejects a second job while busy', async () => {
  const f = fixture()
  const dataDir = join(f.root, 'jobs2')
  let resolveHealth
  const healthGate = new Promise(resolve => { resolveHealth = resolve })
  const hung = createCaptionHost({
    modelsDir: join(f.root, 'models'),
    llamaServer: join(f.root, 'llama-server.exe'),
    dataDir,
    prepare: async () => {},
    leaseValid: () => true,
    spawnProcess: () => {
      const child = new EventEmitter()
      child.pid = 99
      child.stdout = new PassThrough()
      child.stderr = new PassThrough()
      child.kill = () => { setImmediate(() => child.emit('close', 1)); return true }
      child.exitCode = null
      return child
    },
    fetchImpl: async (url) => {
      if (String(url).endsWith('/health')) {
        await healthGate
        return { ok: true }
      }
      return { ok: true, json: async () => ({ choices: [{ message: { content: 'x' } }] }) }
    }
  })
  mkdirSync(join(dataDir, 'caption-job-aaaaaa'), { recursive: true })
  writeFileSync(join(dataDir, 'caption-job-aaaaaa', 'input.png'), Buffer.alloc(64, 1))
  const first = hung.start({ id: 'caption-job-aaaaaa', style: 'tags' }, 'lease')
  await new Promise(r => setTimeout(r, 20))
  await assert.rejects(hung.start({ id: 'caption-job-bbbbbb', style: 'tags' }, 'lease'), /already running/)
  await hung.cancel('caption-job-aaaaaa')
  resolveHealth()
  await assert.rejects(first)
})

test('POST /caption sync shape is documented on host module', () => {
  assert.equal(typeof createCaptionHost({ prepare: async () => {}, leaseValid: () => true }).captionOnce, 'function')
})