import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { AgentId } from '../../core/agent.ts'
import type { Modality, ModelEntry, ModelRegistry, ProviderEntry } from '../../core/model-registry.ts'
import type { NormalizedBlock } from '../../core/packet.ts'
import type { IRRequest } from '../translate/index.ts'
import type { AttemptResult } from './exec.ts'
import type { ResolvedMember, ResolvedToolModel } from './resolve.ts'

const mockReg = vi.hoisted(() => ({
  registry: { version: 1, providers: [], models: [] } as ModelRegistry,
}))

vi.mock('../routing/load.ts', () => ({
  getModelRegistry: () => mockReg.registry,
  getRoutingPolicy: () => ({ version: 2, strategies: [], agents: {} }),
  getSecrets: () => ({ version: 1, secrets: {} }),
}))

vi.mock('./exec.ts', async (importOriginal) => {
  const actual = await importOriginal<typeof import('./exec.ts')>()
  return { ...actual, execAttempt: vi.fn() }
})

const { execAttempt } = await import('./exec.ts')
const { assessVision, countImages, executeVisionLoop, SYNTHETIC_TOOLS } = await import(
  './vision-loop.ts'
)

const mockExec = vi.mocked(execAttempt)

// ── fixtures ──────────────────────────────────────────────────────

const CTX = { agent: 'claude-code' as AgentId, reqHeaders: new Headers() }

function model(id: string, input: Modality[]): ModelEntry {
  return {
    id,
    providerId: `p-${id}`,
    label: id,
    api: 'anthropic-messages',
    input,
    origin: 'manual',
  }
}

function provider(id: string): ProviderEntry {
  return {
    id: `p-${id}`,
    label: id,
    baseUrl: `https://up/${id}`,
    api: 'anthropic-messages',
    auth: { kind: 'passthrough' },
    origin: 'manual',
  }
}

function member(m: ModelEntry): ResolvedMember {
  return { model: m, provider: provider(m.id), api: 'anthropic-messages', switchOn: [] }
}

function visionTool(m: ModelEntry): ResolvedToolModel {
  return { kind: 'vision', model: m, provider: provider(m.id), api: 'anthropic-messages' }
}

function okResult(modelId: string, blocks: NormalizedBlock[]): AttemptResult {
  return {
    ok: true,
    status: 200,
    json: {},
    ir: { model: modelId, blocks, stopReason: 'end_turn', usage: { in: 10, out: 5 } },
    durMs: 1,
    startedAtWall: Date.now(),
    upstreamUrl: `https://up/${modelId}`,
  }
}

function errResult(modelId: string, status: number, errorText: string): AttemptResult {
  return {
    ok: false,
    status,
    errorText,
    durMs: 1,
    startedAtWall: Date.now(),
    upstreamUrl: `https://up/${modelId}`,
  }
}

function text(t: string): NormalizedBlock {
  return { type: 'text', text: t }
}

function viewImage(id: string, imageId: number, question: string): NormalizedBlock {
  return { type: 'tool_use', id, name: 'view_image', input: { image_id: imageId, question } }
}

/** An Anthropic-Messages request carrying one inline image (data `IMGDATA`). */
function reqWithImage(): Record<string, unknown> {
  return {
    model: 'client-model',
    max_tokens: 1024,
    messages: [
      {
        role: 'user',
        content: [
          { type: 'text', text: 'What color is this?' },
          { type: 'image', source: { type: 'base64', media_type: 'image/png', data: 'IMGDATA' } },
        ],
      },
    ],
  }
}

// Each upstream call is scripted per model id and consumed in order.
const scripts = new Map<string, AttemptResult[]>()

beforeEach(() => {
  scripts.clear()
  mockReg.registry = { version: 1, providers: [], models: [] }
  mockExec.mockReset()
  mockExec.mockImplementation(async (m: ResolvedMember) => {
    const next = scripts.get(m.model.id)?.shift()
    if (!next) throw new Error(`vision-loop test: no scripted result for ${m.model.id}`)
    return next
  })
})

// ── countImages ───────────────────────────────────────────────────

describe('countImages', () => {
  it('counts top-level image blocks', () => {
    const ir: IRRequest = {
      model: 'm',
      stream: false,
      messages: [
        {
          role: 'user',
          content: [
            { type: 'text', text: 'hi' },
            { type: 'image', source: { kind: 'base64', mediaType: 'image/png', data: 'X' } },
          ],
        },
      ],
    }
    expect(countImages(ir)).toBe(1)
  })

  it('counts images nested in tool_result content', () => {
    const ir: IRRequest = {
      model: 'm',
      stream: false,
      messages: [
        {
          role: 'user',
          content: [
            {
              type: 'tool_result',
              toolUseId: 't1',
              content: [
                { type: 'image', source: { kind: 'base64', mediaType: 'image/png', data: 'A' } },
                { type: 'image', source: { kind: 'url', url: 'https://x/y.png' } },
              ],
            },
          ],
        },
      ],
    }
    expect(countImages(ir)).toBe(2)
  })

  it('returns 0 when there are no images', () => {
    const ir: IRRequest = {
      model: 'm',
      stream: false,
      messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }],
    }
    expect(countImages(ir)).toBe(0)
  })
})

// ── SYNTHETIC_TOOLS ───────────────────────────────────────────────

describe('SYNTHETIC_TOOLS', () => {
  it('contains view_image', () => {
    expect(SYNTHETIC_TOOLS.has('view_image')).toBe(true)
  })
})

// ── assessVision ──────────────────────────────────────────────────

describe('assessVision', () => {
  const visionToolModel = visionTool(model('vision-companion', ['text', 'image']))

  it('returns none when the strategy has no vision tool-model', () => {
    const textMember = member(model('text-only', ['text']))
    expect(
      assessVision(textMember, undefined, 'anthropic-messages', reqWithImage()).kind,
    ).toBe('none')
  })

  it('returns none when the member can already see images', () => {
    const textMember = member(model('multimodal', ['text', 'image']))
    expect(
      assessVision(textMember, visionToolModel, 'anthropic-messages', reqWithImage()).kind,
    ).toBe('none')
  })

  it('returns none when a text-only member receives no images', () => {
    const textMember = member(model('text-only', ['text']))
    const req = { model: 'm', max_tokens: 64, messages: [{ role: 'user', content: 'hi' }] }
    expect(assessVision(textMember, visionToolModel, 'anthropic-messages', req).kind).toBe('none')
  })

  it('returns a loop with both members when a tool-model and images are present', () => {
    const textMember = member(model('text-only', ['text']))
    const result = assessVision(
      textMember,
      visionToolModel,
      'anthropic-messages',
      reqWithImage(),
    )
    expect(result.kind).toBe('loop')
    if (result.kind !== 'loop') return
    expect(result.textMember.model.id).toBe('text-only')
    expect(result.visionMember.model.id).toBe('vision-companion')
  })
})

// ── executeVisionLoop ─────────────────────────────────────────────

describe('executeVisionLoop', () => {
  const textModel = member(model('text-model', ['text']))
  const visionModel = member(model('vision-model', ['text', 'image']))

  function run(maxIterations: number) {
    return executeVisionLoop({
      clientApi: 'anthropic-messages',
      clientRequest: reqWithImage(),
      textMember: textModel,
      visionMember: visionModel,
      ctx: CTX,
      maxIterations,
    })
  }

  it('substitutes placeholders, services view_image, and returns the final answer', async () => {
    scripts.set('text-model', [
      okResult('text-model', [viewImage('tu_1', 1, 'What color?')]),
      okResult('text-model', [text('The image shows a red square.')]),
    ])
    scripts.set('vision-model', [okResult('vision-model', [text('Red')])])

    const outcome = await run(6)

    expect(outcome.ok).toBe(true)
    expect(outcome.risk).toEqual([])
    expect(outcome.attempts.map((a) => a.role)).toEqual(['text-turn', 'vision', 'text-turn'])
    expect(outcome.finalIR.blocks).toEqual([text('The image shows a red square.')])

    // The text model sees a placeholder and the injected tool — never the bytes.
    const turn1 = JSON.stringify(mockExec.mock.calls[0]?.[2] ?? {})
    expect(turn1).toContain('[Image #1]')
    expect(turn1).toContain('view_image')
    expect(turn1).not.toContain('IMGDATA')

    // The vision model is the one that receives the actual image.
    expect(JSON.stringify(mockExec.mock.calls[1]?.[2] ?? {})).toContain('IMGDATA')

    // The vision answer is fed back into the next text turn as a tool_result.
    expect(JSON.stringify(mockExec.mock.calls[2]?.[2] ?? {})).toContain('Red')
  })

  it('returns an error tool_result and makes no vision call for a bad image id', async () => {
    scripts.set('text-model', [
      okResult('text-model', [viewImage('tu_1', 99, 'What color?')]),
      okResult('text-model', [text('I could not inspect the image.')]),
    ])

    const outcome = await run(6)

    expect(outcome.ok).toBe(true)
    expect(mockExec).toHaveBeenCalledTimes(2)
    expect(outcome.attempts.map((a) => a.role)).toEqual(['text-turn', 'text-turn'])
    const fed = JSON.stringify(mockExec.mock.calls[1]?.[2] ?? {})
    expect(fed).toContain('No image with id 99')
    expect(fed).toContain('is_error')
  })

  it('exhausts the loop and falls back when view_image calls never stop', async () => {
    scripts.set('text-model', [
      okResult('text-model', [viewImage('tu_1', 1, 'q1')]),
      okResult('text-model', [viewImage('tu_2', 1, 'q2')]),
    ])
    scripts.set('vision-model', [okResult('vision-model', [text('Red')])])

    const outcome = await run(2)

    expect(outcome.ok).toBe(true)
    expect(outcome.risk.map((r) => r.tag)).toContain('vision:loop-exhausted')
    expect(outcome.finalIR.blocks).toEqual([text('Unable to complete image analysis.')])
  })

  it('ends with ok:false and a high-severity risk when a text turn fails', async () => {
    scripts.set('text-model', [errResult('text-model', 502, 'upstream exploded')])

    const outcome = await run(6)

    expect(outcome.ok).toBe(false)
    const failRisk = outcome.risk.find((r) => r.tag === 'vision:text-failed')
    expect(failRisk?.severity).toBe('high')
    expect(outcome.finalIR.blocks).toEqual([text('Unable to complete image analysis.')])
  })
})