import { beforeEach, describe, expect, it, vi } from 'vitest'
import type { AgentStatus, AiAgentMessage } from './aiAgentConversation'
import type { AiModelDefinition, AiModelProvider, AiTarget } from './aiTargets'

const {
  buildAgentSystemPromptMock,
  createStreamCallbacksMock,
  formatMessageWithHistoryMock,
  hydrateNoteReferencesMock,
  nextMessageIdMock,
  streamAiAgentMock,
  streamAiModelMock,
  trackEventMock,
  trimHistoryMock,
} = vi.hoisted(() => ({
  buildAgentSystemPromptMock: vi.fn(() => 'SYSTEM'),
  createStreamCallbacksMock: vi.fn(() => ({ stream: 'callbacks' })),
  formatMessageWithHistoryMock: vi.fn((_history: unknown, prompt: string) => `formatted:${prompt}`),
  hydrateNoteReferencesMock: vi.fn(async (references: unknown) => references),
  nextMessageIdMock: vi.fn(),
  streamAiAgentMock: vi.fn(async () => {}),
  streamAiModelMock: vi.fn(async () => {}),
  trackEventMock: vi.fn(),
  trimHistoryMock: vi.fn((history: unknown) => history),
}))

vi.mock('../utils/ai-agent', () => ({
  buildAgentSystemPrompt: buildAgentSystemPromptMock,
}))

vi.mock('../utils/ai-chat', () => ({
  MAX_HISTORY_TOKENS: 100_000,
  formatMessageWithHistory: formatMessageWithHistoryMock,
  nextMessageId: nextMessageIdMock,
  trimHistory: trimHistoryMock,
}))

vi.mock('./aiAgentStreamCallbacks', () => ({
  createStreamCallbacks: createStreamCallbacksMock,
}))

vi.mock('../utils/streamAiAgent', () => ({
  streamAiAgent: streamAiAgentMock,
}))

vi.mock('../utils/streamAiModel', () => ({
  streamAiModel: streamAiModelMock,
}))

vi.mock('../utils/ai-reference-content', () => ({
  hydrateNoteReferences: hydrateNoteReferencesMock,
}))

vi.mock('./telemetry', () => ({
  trackEvent: trackEventMock,
}))

import {
  clearAgentConversation,
  sendAgentMessage,
  stopAgentMessage,
  type AiAgentSessionRuntime,
} from './aiAgentSession'

function createRuntime(
  initialMessages: AiAgentMessage[] = [],
  initialStatus: AgentStatus = 'idle',
) {
  let messages = initialMessages
  let status = initialStatus

  const messagesRef = { current: messages }
  const statusRef = { current: status }

  const setMessages = vi.fn((next: AiAgentMessage[] | ((current: AiAgentMessage[]) => AiAgentMessage[])) => {
    messages = typeof next === 'function' ? next(messages) : next
    messagesRef.current = messages
  })
  const setStatus = vi.fn((next: AgentStatus | ((current: AgentStatus) => AgentStatus)) => {
    status = typeof next === 'function' ? next(status) : next
    statusRef.current = status
  })

  const runtime: AiAgentSessionRuntime = {
    setMessages,
    setStatus,
    abortRef: { current: { aborted: true } },
    responseAccRef: { current: 'stale response' },
    fileCallbacksRef: { current: { onVaultChanged: vi.fn() } },
    toolInputMapRef: { current: new Map([['stale-tool', { tool: 'Write', input: '{"path":"/stale.md"}' }]]) },
    messagesRef,
    statusRef,
  }

  return {
    runtime,
    getMessages: () => messages,
    getStatus: () => status,
  }
}

type RuntimeFixture = ReturnType<typeof createRuntime>

const completedHistory: AiAgentMessage = {
  id: 'msg-1',
  userMessage: 'Previous question',
  actions: [],
  response: 'Previous answer',
}
const streamingHistory: AiAgentMessage = {
  id: 'msg-2',
  userMessage: 'Ignored streaming question',
  actions: [],
  isStreaming: true,
}
const expectedChatHistory = [
  { role: 'user', content: 'Previous question', id: 'msg-1' },
  { role: 'assistant', content: 'Previous answer', id: 'msg-1-resp' },
]
const apiModelProvider: AiModelProvider = {
  id: 'openai',
  name: 'OpenAI',
  kind: 'open_ai',
  base_url: 'https://api.openai.com/v1',
  api_key_storage: 'local_file',
  api_key_env_var: null,
  models: [],
}
const apiModel: AiModelDefinition = {
  id: 'gpt-5-nano',
  display_name: 'GPT-5 nano',
  context_window: null,
  max_output_tokens: null,
  capabilities: {
    streaming: false,
    tools: false,
    vision: false,
    json_mode: false,
    reasoning: false,
  },
}
const apiTarget: AiTarget = {
  kind: 'api_model',
  provider: apiModelProvider,
  model: apiModel,
  id: 'model:openai/gpt-5-nano',
  label: 'OpenAI · GPT-5 nano',
  shortLabel: 'GPT-5 nano',
}

function expectStreamingRuntimeState(session: RuntimeFixture): void {
  expect(session.runtime.abortRef.current.aborted).toBe(false)
  expect(session.runtime.abortRef.current.controller).toBeInstanceOf(AbortController)
  expect(session.runtime.responseAccRef.current).toBe('')
  expect(session.runtime.toolInputMapRef.current.size).toBe(0)
  expect(session.getStatus()).toBe('thinking')
  expect(session.getMessages().at(-1)).toEqual({
    userMessage: 'Latest question',
    references: [{ path: '/vault/ref.md', title: 'Ref' }],
    actions: [],
    isStreaming: true,
    id: 'msg-stream',
  })
}

function expectFormattedHistoryUsed(): void {
  expect(trimHistoryMock).toHaveBeenCalledWith(expectedChatHistory, 100_000)
  expect(formatMessageWithHistoryMock).toHaveBeenCalledWith(
    expectedChatHistory,
    expect.stringContaining('Latest question'),
  )
  expect(formatMessageWithHistoryMock).toHaveBeenCalledWith(
    expectedChatHistory,
    expect.stringContaining('/vault/ref.md'),
  )
}

function expectStreamingRequest(runtime: RuntimeFixture['runtime']): void {
  expect(createStreamCallbacksMock).toHaveBeenCalledWith(expect.objectContaining({
    messageId: 'msg-stream',
    locale: 'it-IT',
    vaultPath: '/vault',
    setMessages: runtime.setMessages,
    setStatus: runtime.setStatus,
  }))
  expect(streamAiAgentMock).toHaveBeenCalledWith(expect.objectContaining({
    agent: 'codex',
    message: expect.stringContaining('formatted:Latest question'),
    systemPrompt: 'SYSTEM',
    vaultPath: '/vault',
    permissionMode: 'power_user',
    callbacks: { stream: 'callbacks' },
    signal: expect.any(AbortSignal),
  }))
}

function expectApiModelStreamingRequest(runtime: RuntimeFixture['runtime']): void {
  expect(createStreamCallbacksMock).toHaveBeenCalledWith(expect.objectContaining({
    messageId: 'msg-stream',
    vaultPath: '/vault',
    setMessages: runtime.setMessages,
    setStatus: runtime.setStatus,
  }))
  expect(streamAiModelMock).toHaveBeenCalledWith({
    provider: apiModelProvider,
    model: apiModel,
    message: expect.stringContaining('formatted:Latest question'),
    systemPrompt: 'SYSTEM',
    vaultPath: '/vault',
    vaultPaths: ['/vault', '/team-vault'],
    callbacks: { stream: 'callbacks' },
  })
}

describe('aiAgentSession', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    buildAgentSystemPromptMock.mockReturnValue('SYSTEM')
    createStreamCallbacksMock.mockReturnValue({ stream: 'callbacks' })
    formatMessageWithHistoryMock.mockImplementation((_history: unknown, prompt: string) => `formatted:${prompt}`)
    trimHistoryMock.mockImplementation((history: unknown) => history)
    streamAiAgentMock.mockResolvedValue(undefined)
    hydrateNoteReferencesMock.mockImplementation(async (references: unknown) => references)
    trackEventMock.mockClear()
  })

  async function expectLocalResponse(options: {
    messageId: string
    context: {
      agent: 'claude_code' | 'codex' | 'copilot' | 'opencode' | 'pi' | 'antigravity'
      ready: boolean
      vaultPath: string
      permissionMode: 'safe' | 'power_user'
    }
    prompt: { text: string; references?: [] }
    reason: 'agent_unavailable' | 'missing_vault'
    response: string
  }) {
    nextMessageIdMock.mockReturnValue(options.messageId)
    const { runtime, getMessages } = createRuntime()

    await sendAgentMessage({
      runtime,
      context: options.context,
      prompt: options.prompt,
    })

    expect(getMessages()).toEqual([
      {
        userMessage: options.prompt.text,
        references: undefined,
        actions: [],
        response: options.response,
        id: options.messageId,
      },
    ])
    expect(streamAiAgentMock).not.toHaveBeenCalled()
    expect(trackEventMock).toHaveBeenCalledWith('ai_agent_message_blocked', {
      agent: options.context.agent,
      reason: options.reason,
    })
  }

  it('ignores blank prompts and busy runtimes', async () => {
    const idleRuntime = createRuntime()
    await sendAgentMessage({
      runtime: idleRuntime.runtime,
      context: { agent: 'codex', ready: true, vaultPath: '/vault', permissionMode: 'safe' },
      prompt: { text: '   ' },
    })

    const busyRuntime = createRuntime([], 'thinking')
    await sendAgentMessage({
      runtime: busyRuntime.runtime,
      context: { agent: 'codex', ready: true, vaultPath: '/vault', permissionMode: 'safe' },
      prompt: { text: 'Question' },
    })

    expect(idleRuntime.getMessages()).toEqual([])
    expect(busyRuntime.getMessages()).toEqual([])
    expect(streamAiAgentMock).not.toHaveBeenCalled()
  })

  it('appends local fallback responses when the session cannot stream', async () => {
    const fallbackCases = [
      {
        messageId: 'msg-local',
        context: { agent: 'codex', ready: true, vaultPath: '', permissionMode: 'safe' },
        prompt: { text: 'Open a note' },
        reason: 'missing_vault',
        response: 'No vault loaded. Open a vault first.',
      },
      {
        messageId: 'msg-missing',
        context: { agent: 'codex', ready: false, vaultPath: '/vault', permissionMode: 'safe' },
        prompt: { text: 'Open a note', references: [] },
        reason: 'agent_unavailable',
        response:
          'Codex is not available on this machine. Install it or switch the default AI agent in Settings.',
      },
    ] as const

    for (const fallbackCase of fallbackCases) {
      await expectLocalResponse(fallbackCase)
    }
  })

  it('starts a streaming session with formatted history and fresh refs', async () => {
    nextMessageIdMock.mockReturnValue('msg-stream')
    const session = createRuntime([
      completedHistory,
      streamingHistory,
    ])

    await sendAgentMessage({
      runtime: session.runtime,
      context: {
        agent: 'codex',
        locale: 'it-IT',
        ready: true,
        vaultPath: '/vault',
        permissionMode: 'power_user',
        systemPromptOverride: 'OVERRIDE',
      },
      prompt: {
        text: '  Latest question  ',
        references: [{ path: '/vault/ref.md', title: 'Ref' }],
      },
    })

    expectStreamingRuntimeState(session)
    expect(hydrateNoteReferencesMock).toHaveBeenCalledWith([{ path: '/vault/ref.md', title: 'Ref' }])
    expectFormattedHistoryUsed()
    expect(buildAgentSystemPromptMock).toHaveBeenCalledWith({
      agent: 'codex',
      permissionMode: 'power_user',
      vaultContext: 'OVERRIDE',
    })
    expectStreamingRequest(session.runtime)
    expect(trackEventMock).toHaveBeenCalledWith('ai_agent_message_sent', {
      agent: 'codex',
      permission_mode: 'power_user',
      has_context: 1,
      reference_count: 1,
      history_message_count: 1,
    })
  })

  it('passes vault roots to api model streams for native note tools', async () => {
    nextMessageIdMock.mockReturnValue('msg-stream')
    const session = createRuntime([
      completedHistory,
      streamingHistory,
    ])

    await sendAgentMessage({
      runtime: session.runtime,
      context: {
        agent: 'codex',
        target: apiTarget,
        ready: true,
        vaultPath: '/vault',
        vaultPaths: ['/vault', '/team-vault'],
        permissionMode: 'safe',
      },
      prompt: {
        text: '  Latest question  ',
        references: [{ path: '/vault/ref.md', title: 'Ref' }],
      },
    })

    expectStreamingRuntimeState(session)
    expectFormattedHistoryUsed()
    expectApiModelStreamingRequest(session.runtime)
    expect(streamAiAgentMock).not.toHaveBeenCalled()
  })

  it('clears the conversation and resets runtime refs', () => {
    const { runtime } = createRuntime([
      { id: 'msg-1', userMessage: 'Question', actions: [] },
    ], 'done')

    clearAgentConversation(runtime)

    expect(runtime.abortRef.current.aborted).toBe(true)
    expect(runtime.responseAccRef.current).toBe('')
    expect(runtime.toolInputMapRef.current.size).toBe(0)
    expect(runtime.setMessages).toHaveBeenCalledWith([])
    expect(runtime.setStatus).toHaveBeenCalledWith('idle')
  })

  it('stops the active stream and marks the streaming message as stopped', async () => {
    nextMessageIdMock.mockReturnValue('msg-stream')
    const session = createRuntime()
    let streamSignal: AbortSignal | undefined
    streamAiAgentMock.mockImplementation(async ({ signal }: { signal?: AbortSignal }) => new Promise<void>((resolve) => {
      streamSignal = signal
      signal?.addEventListener('abort', () => resolve(), { once: true })
    }))

    const pending = sendAgentMessage({
      runtime: session.runtime,
      context: {
        agent: 'codex',
        ready: true,
        vaultPath: '/vault',
        permissionMode: 'safe',
      },
      prompt: { text: '  Latest question  ' },
    })
    await Promise.resolve()
    await Promise.resolve()

    stopAgentMessage(session.runtime, { agent: 'codex', locale: 'en' })
    await pending

    expect(streamSignal?.aborted).toBe(true)
    expect(session.runtime.abortRef.current.aborted).toBe(true)
    expect(session.getStatus()).toBe('idle')
    expect(session.getMessages()).toEqual([{
      userMessage: 'Latest question',
      actions: [],
      isStreaming: false,
      reasoningDone: true,
      response: 'Stopped.',
      id: 'msg-stream',
    }])
    expect(trackEventMock).toHaveBeenCalledWith('ai_agent_response_stopped', {
      agent: 'codex',
      had_partial_response: 0,
      tool_count: 0,
    })
  })
})
