import { describe, expect, it, vi } from 'vitest'
import { createDimSdkState } from '../src/dimsdk/core/store/state.ts'
import { createNativeStreamRunner } from '../src/dimsdk/core/store/native-stream-runner.ts'

function createRunnerHarness(input?: {
  run?: () => AsyncIterable<unknown>
  applyLastRequestBudget?: (budget: unknown) => void
  applyUsageFromSession?: (usage: unknown) => void
}) {
  const { state } = createDimSdkState({
    context: { cwd: process.cwd(), env: {} } as any,
    agent: { getMode: () => 'agent' } as any,
  })
  state.config.context_window = 4096
  state.config.model = 'gpt-5'
  state.messages.push({
    id: 'user-1',
    role: 'user',
    content: 'hello',
    createdAtMs: Date.now(),
  } as any)

  const ensureSession = vi.fn(async () => ({
    id: 'session-1',
    run: input?.run ?? (async function* () {}),
  }))

  const applyLastRequestBudget = input?.applyLastRequestBudget ?? vi.fn()
  const applyUsageFromSession = input?.applyUsageFromSession ?? vi.fn()

  const runner = createNativeStreamRunner({
    state,
    agent: {
      ensureSession,
      getMode: () => 'agent',
    } as any,
    id: (prefix: string) => `${prefix}-1`,
    pushMessage: (message) => {
      state.messages.push(message)
    },
    updateAssistant: () => {},
    setAssistantStatus: () => {},
    setAssistantThinking: () => {},
    startNewThinkingBlock: () => {},
    endThinkingBlock: () => {},
    appendAssistantMarkdown: () => {},
    pushAssistantMarkdownPart: () => {},
    appendInfoPart: () => {},
    upsertSubagentInfo: () => {},
    tickSubagentFrame: () => {},
    upsertToolCallDelta: () => {},
    upsertToolCallEnd: () => {},
    upsertToolOutputStart: () => {},
    appendToolOutputDelta: () => {},
    insertApprovalPart: () => {},
    attachToolResultPart: () => {},
    scheduleToolResultHighlight: () => {},
    upsertChecklistPartFromDimSdk: () => {},
    applyLastRequestBudget,
    applyUsageFromSession,
    ensureErrorVisibleInTranscript: () => {},
  })

  return {
    applyLastRequestBudget,
    applyUsageFromSession,
    ensureSession,
    runner,
    state,
  }
}

describe('native stream runner', () => {
  it('passes the configured context window into runtime config without throwing', async () => {
    const { ensureSession, runner } = createRunnerHarness()

    const result = await runner.runAgentRound()

    expect(result.success).toBe(true)
    expect(ensureSession).toHaveBeenCalledWith(expect.objectContaining({
      runtimeConfig: expect.objectContaining({
        maxInputTokens: 4096,
      }),
    }))
  })

  it('applies request budget updates from usage.updated events', async () => {
    const applyLastRequestBudget = vi.fn()
    const applyUsageFromSession = vi.fn()
    const { runner } = createRunnerHarness({
      applyLastRequestBudget,
      applyUsageFromSession,
      run: async function* () {
        yield {
          type: 'usage.updated',
          lastRequestBudget: {
            autoBudgetOutcome: 'budget-unknown',
            estimatedInputTokens: 1800,
            updatedAtMs: 1712966400000,
            usedAutoCompaction: false,
          },
          usage: {
            promptTokens: 1800,
            completionTokens: 90,
            totalTokens: 1890,
          },
        } as any
        yield {
          type: 'run.finished',
          status: 'completed',
          reasoningDiagnostics: {},
        } as any
      },
    })

    const result = await runner.runAgentRound()

    expect(result.success).toBe(true)
    expect(applyUsageFromSession).toHaveBeenCalledWith({
      promptTokens: 1800,
      completionTokens: 90,
      totalTokens: 1890,
    })
    expect(applyLastRequestBudget).toHaveBeenCalledWith(
      expect.objectContaining({
        autoBudgetOutcome: 'budget-unknown',
        estimatedInputTokens: 1800,
      }),
    )
  })

  it('normalizes legacy run.budget events into the request-budget snapshot shape', async () => {
    const applyLastRequestBudget = vi.fn()
    const { runner } = createRunnerHarness({
      applyLastRequestBudget,
      run: async function* () {
        yield {
          type: 'run.budget',
          runId: 'run-1',
          requestId: 'req-1',
          estimatedInputTokens: 1331,
          thresholdTokens: 3200,
          maxInputTokens: 3200,
          triggerRatio: 1,
          usedAutoCompaction: false,
        } as any
        yield {
          type: 'run.finished',
          status: 'completed',
          reasoningDiagnostics: {},
        } as any
      },
    })

    const result = await runner.runAgentRound()

    expect(result.success).toBe(true)
    expect(applyLastRequestBudget).toHaveBeenCalledWith(
      expect.objectContaining({
        autoBudgetOutcome: 'not-needed',
        estimatedInputTokens: 1331,
        inputBudgetTokens: 3200,
        thresholdTokens: 3200,
        maxInputTokens: 3200,
      }),
    )
  })
})
