import { messageTable } from '@data/db/schemas/message'
import { topicTable } from '@data/db/schemas/topic'
import { userModelTable } from '@data/db/schemas/userModel'
import { userProviderTable } from '@data/db/schemas/userProvider'
import { messageService } from '@data/services/MessageService'
import { topicService } from '@data/services/TopicService'
import { generateOrderKeySequence } from '@data/services/utils/orderKey'
import type { AiStreamOpenRequest } from '@shared/ai/transport'
import { createUniqueModelId } from '@shared/data/types/model'
import { setupTestDatabase, withRoot } from '@test-helpers/db'
import { beforeEach, describe, expect, it, vi } from 'vitest'

import { startAiChildTurnSpan } from '../../../observability'
import { PersistenceListener } from '../../listeners/PersistenceListener'
import type { StreamListener } from '../../types'
import type { MainSteerContinuationRequest } from '../dispatch'
import { resolveModels, resolvePersistentSiblingsGroupId } from '../modelResolution'

// Stub model resolution + tracing so the test drives the REAL DB history path
// (`createUserMessageWithPlaceholders` → `getPathToNode`) without provider/model
// resolution machinery. The history is what we assert on.
const MODEL_ID = createUniqueModelId('openai', 'gpt-4o')
vi.mock('../modelResolution', () => ({
  resolveAssistantModelId: vi.fn(async () => ({ assistantId: undefined, defaultModelId: MODEL_ID })),
  resolveModels: vi.fn(async () => [{ id: MODEL_ID, name: 'GPT-4o', providerId: 'openai', apiModelId: 'gpt-4o' }]),
  resolvePersistentSiblingsGroupId: vi.fn(async () => 1)
}))

vi.mock('../../../observability', () => ({
  startAiChildTurnSpan: vi.fn(() => ({ rootSpan: { end: vi.fn() }, traceId: 'trace-1' })),
  deriveRootSpanId: vi.fn(() => '1'.repeat(16)),
  applyTurnInputAttributes: vi.fn()
}))

const { PersistentChatContextProvider } = await import('../PersistentChatContextProvider')

function makeSubscriber(): StreamListener {
  return { id: 'wc:1', onChunk: vi.fn(), onDone: vi.fn(), onPaused: vi.fn(), onError: vi.fn(), isAlive: () => true }
}

/** Flatten a history message to `{ role, text }` for order-sensitive assertions. */
function flatten(messages: { role: string; parts: Array<{ type: string; text?: string }> }[]) {
  return messages.map((m) => ({
    role: m.role,
    text: m.parts
      .filter((p) => p.type === 'text')
      .map((p) => p.text ?? '')
      .join('')
  }))
}

describe('PersistentChatContextProvider — steer continuation history', () => {
  const dbh = setupTestDatabase()
  const provider = new PersistentChatContextProvider()

  // The text a prior turn produced before it yielded to the steer; the steer continuation's
  // history must include it (it was persisted on the assistant row by the normal terminal path).
  const PARTIAL = 'partial answer so far'

  beforeEach(async () => {
    const [providerKey, modelKey] = generateOrderKeySequence(2)
    await dbh.db.insert(userProviderTable).values({ providerId: 'openai', name: 'OpenAI', orderKey: providerKey })
    await dbh.db.insert(userModelTable).values({
      id: MODEL_ID,
      providerId: 'openai',
      modelId: 'gpt-4o',
      presetModelId: 'gpt-4o',
      name: 'GPT-4o',
      isEnabled: true,
      isHidden: false,
      orderKey: modelKey
    })

    await dbh.db.insert(topicTable).values({ id: 'topic-1', activeNodeId: 'a1', orderKey: 'a0' })
    await dbh.db.insert(messageTable).values(
      withRoot('topic-1', [
        {
          id: 'u1',
          parentId: null,
          topicId: 'topic-1',
          role: 'user',
          data: { parts: [{ type: 'text', text: 'first question' }] },
          status: 'success',
          siblingsGroupId: 0,
          createdAt: 100,
          updatedAt: 100
        },
        {
          id: 'a1',
          parentId: 'u1',
          topicId: 'topic-1',
          role: 'assistant',
          data: { parts: [{ type: 'text', text: PARTIAL }] },
          status: 'paused',
          siblingsGroupId: 1,
          modelId: MODEL_ID,
          createdAt: 200,
          updatedAt: 200
        }
      ])
    )
  })

  it('rebuilds a prompt that carries the paused partial when the new turn anchors on the paused row', async () => {
    // Steering: renderer's `activeNodeId` (the streaming/paused assistant row) is sent as
    // `parentAnchorId`, so the new user message is parented on the paused row.
    const prepared = await provider.prepareDispatch(
      makeSubscriber(),
      {
        trigger: 'submit-message',
        topicId: 'topic-1',
        parentAnchorId: 'a1',
        userMessageParts: [{ type: 'text', text: 'actually, change direction' }]
      } as AiStreamOpenRequest,
      { hasLiveStream: false }
    )

    const history = prepared.models[0].request.messages
    expect(history).toBeDefined()
    expect(flatten(history!)).toEqual([
      { role: 'user', text: 'first question' },
      // The paused partial survives into the rebuilt prompt — this is the B4 efficacy guarantee.
      { role: 'assistant', text: PARTIAL },
      { role: 'user', text: 'actually, change direction' }
    ])
  })

  it('steer-continuation: opens an assistant turn under the steer user row with a reminder-wrapped prompt', async () => {
    // u1 → a1 → u2, where u2 is the steer the user sent mid-turn (child of the assistant row).
    await dbh.db.insert(messageTable).values({
      id: 'u2',
      parentId: 'a1',
      topicId: 'topic-1',
      role: 'user',
      data: { parts: [{ type: 'text', text: 'actually do X instead' }] },
      status: 'success',
      siblingsGroupId: 0,
      modelId: MODEL_ID,
      createdAt: 300,
      updatedAt: 300
    })

    const prepared = await provider.prepareDispatch(
      makeSubscriber(),
      { trigger: 'steer-continuation', topicId: 'topic-1', userMessageId: 'u2' } as MainSteerContinuationRequest,
      { hasLiveStream: false }
    )

    // A fresh assistant placeholder is created under u2 — no new user row.
    const children = await messageService.getChildrenByParentId('u2')
    expect(children).toHaveLength(1)
    expect(children[0]).toMatchObject({ role: 'assistant', status: 'pending' })

    // The persisted steer row is untouched (only the model-facing copy is wrapped).
    const u2 = await messageService.getById('u2')
    expect(flatten([{ role: 'user', parts: u2.data.parts ?? [] }])[0].text).toBe('actually do X instead')

    // History is the full path; only the trailing steer message is system-reminder wrapped.
    const history = prepared.models[0].request.messages!
    expect(history).toHaveLength(3)
    expect(flatten([history[0]])[0]).toEqual({ role: 'user', text: 'first question' })
    expect(flatten([history[1]])[0]).toEqual({ role: 'assistant', text: PARTIAL })
    const lastText = flatten([history[2]])[0].text
    expect(history[2].role).toBe('user')
    expect(lastText).toContain('<system-reminder>')
    expect(lastText).toContain('actually do X instead')
    expect(lastText).toContain('Please address this message and continue with your tasks.')
  })

  it('drops the paused partial when the new turn does not anchor on it (precondition is necessary)', async () => {
    // Counter-case: anchoring on the prior user message (not the paused assistant row) rebuilds
    // a prompt WITHOUT the partial — proving the efficacy hinges on `parentAnchorId` = paused row.
    const prepared = await provider.prepareDispatch(
      makeSubscriber(),
      {
        trigger: 'submit-message',
        topicId: 'topic-1',
        parentAnchorId: 'u1',
        userMessageParts: [{ type: 'text', text: 'retry from before' }]
      } as AiStreamOpenRequest,
      { hasLiveStream: false }
    )

    expect(flatten(prepared.models[0].request.messages!)).toEqual([
      { role: 'user', text: 'first question' },
      { role: 'user', text: 'retry from before' }
    ])
  })

  it('fans out @-mentioned siblings: shared siblingsGroupId, one placeholder per model, aligned placeholders[i]/turnRootSpans[i]', async () => {
    // Two @-mentioned models → two assistant placeholders sharing one siblings group.
    // All placeholders share the container traceId now; assert the per-model row and span
    // line up by index (keyed on modelId) so a fan-out never crosses streams.
    const MODEL_A = createUniqueModelId('openai', 'gpt-4o') // already seeded in beforeEach
    const MODEL_B = createUniqueModelId('anthropic', 'claude-sonnet-4-5')
    // Placeholder rows FK to user_model(id) — seed the second @-mentioned model.
    const [bProviderKey, bModelKey] = generateOrderKeySequence(2)
    await dbh.db
      .insert(userProviderTable)
      .values({ providerId: 'anthropic', name: 'Anthropic', orderKey: bProviderKey })
    await dbh.db.insert(userModelTable).values({
      id: MODEL_B,
      providerId: 'anthropic',
      modelId: 'claude-sonnet-4-5',
      presetModelId: 'claude-sonnet-4-5',
      name: 'Claude Sonnet 4.5',
      isEnabled: true,
      isHidden: false,
      orderKey: bModelKey
    })
    vi.mocked(resolveModels).mockResolvedValueOnce([
      { id: MODEL_A, name: 'GPT-4o', providerId: 'openai', apiModelId: 'gpt-4o' },
      { id: MODEL_B, name: 'Claude Sonnet 4.5', providerId: 'anthropic', apiModelId: 'claude-sonnet-4-5' }
    ] as Awaited<ReturnType<typeof resolveModels>>)
    vi.mocked(resolvePersistentSiblingsGroupId).mockResolvedValueOnce(42)
    // Distinct span per call so index alignment is observable.
    const spanA = { end: vi.fn() }
    const spanB = { end: vi.fn() }
    vi.mocked(startAiChildTurnSpan)
      .mockReturnValueOnce({ rootSpan: spanA } as unknown as ReturnType<typeof startAiChildTurnSpan>)
      .mockReturnValueOnce({ rootSpan: spanB } as unknown as ReturnType<typeof startAiChildTurnSpan>)

    const prepared = await provider.prepareDispatch(
      makeSubscriber(),
      {
        trigger: 'submit-message',
        topicId: 'topic-1',
        parentAnchorId: 'u1',
        mentionedModelIds: [MODEL_A, MODEL_B],
        userMessageParts: [{ type: 'text', text: 'ask both models' }]
      } as AiStreamOpenRequest,
      { hasLiveStream: false }
    )

    // Shared sibling group + multi-model flag.
    expect(prepared.siblingsGroupId).toBe(42)
    expect(prepared.isMultiModel).toBe(true)

    // One execution per model, in mention order, each carrying its own root span.
    expect(prepared.models.map((m) => m.modelId)).toEqual([MODEL_A, MODEL_B])
    expect(prepared.models[0].rootSpan).toBe(spanA)
    expect(prepared.models[1].rootSpan).toBe(spanB)

    // One persisted placeholder per model, both in the shared group, each routed to its
    // own request — placeholders[i]/turnRootSpans[i] alignment proven via per-row modelId.
    const placeholders = await messageService.getChildrenByParentId(prepared.userMessageId!)
    expect(placeholders).toHaveLength(2)
    const byModel = new Map(placeholders.map((p) => [p.modelId, p]))
    const phA = byModel.get(MODEL_A)
    const phB = byModel.get(MODEL_B)
    expect(phA?.modelId).toBe(MODEL_A)
    expect(phB?.modelId).toBe(MODEL_B)
    expect(phA?.siblingsGroupId).toBe(42)
    expect(phB?.siblingsGroupId).toBe(42)
    expect(prepared.models[0].request.messageId).toBe(phA?.id)
    expect(prepared.models[1].request.messageId).toBe(phB?.id)

    // One PersistenceListener per placeholder — no missing/extra/duplicate listener for a fan-out.
    const persistenceListeners = prepared.listeners.filter((l) => l instanceof PersistenceListener)
    expect(persistenceListeners).toHaveLength(2)
    // Each listener is keyed (via its sqlite-backed id `persistence:sqlite:<topicId>:<modelId>`) to the
    // model whose execution carries the matching placeholder messageId — so terminal events route to the
    // right row. modelId order matches the per-model executions, proving listener[i] ↔ placeholder[i].
    expect(persistenceListeners.map((l) => l.id)).toEqual([
      `persistence:sqlite:topic-1:${MODEL_A}`,
      `persistence:sqlite:topic-1:${MODEL_B}`
    ])

    // One trace tree per topic: both models' `ai.turn` spans are parented under the
    // same container trace, and that id is the topic's persisted ensureTraceId.
    const containerTraceId = await topicService.ensureTraceId('topic-1')
    const [callA, callB] = vi.mocked(startAiChildTurnSpan).mock.calls.slice(-2)
    expect(callA[3]).toBe(containerTraceId)
    expect(callB[3]).toBe(containerTraceId)
  })
})

describe('PersistentChatContextProvider — prepareContinueDispatch (resume-after-approval)', () => {
  const dbh = setupTestDatabase()
  const provider = new PersistentChatContextProvider()

  // The anchor's persisted model differs from the test's default model so a
  // reuse failure (resolving the default instead of the anchor) is observable.
  const ANCHOR_MODEL_ID = createUniqueModelId('openai', 'gpt-4o-mini')
  const APPROVAL_ID = 'approval-1'

  beforeEach(async () => {
    vi.clearAllMocks()
    const [providerKey, modelKey, anchorModelKey] = generateOrderKeySequence(3)
    await dbh.db.insert(userProviderTable).values({ providerId: 'openai', name: 'OpenAI', orderKey: providerKey })
    await dbh.db.insert(userModelTable).values([
      {
        id: MODEL_ID,
        providerId: 'openai',
        modelId: 'gpt-4o',
        presetModelId: 'gpt-4o',
        name: 'GPT-4o',
        isEnabled: true,
        isHidden: false,
        orderKey: modelKey
      },
      {
        id: ANCHOR_MODEL_ID,
        providerId: 'openai',
        modelId: 'gpt-4o-mini',
        presetModelId: 'gpt-4o-mini',
        name: 'GPT-4o mini',
        isEnabled: true,
        isHidden: false,
        orderKey: anchorModelKey
      }
    ])

    // topic-2 is a real but empty topic — lets the wrong-topic test reach the belonging
    // guard inside prepareContinueDispatch instead of failing earlier on a missing topic.
    await dbh.db.insert(topicTable).values([
      { id: 'topic-1', activeNodeId: 'a1', orderKey: 'a0' },
      { id: 'topic-2', activeNodeId: null, orderKey: 'a1' }
    ])
    await dbh.db.insert(messageTable).values(
      withRoot('topic-1', [
        {
          id: 'u1',
          parentId: null,
          topicId: 'topic-1',
          role: 'user',
          data: { parts: [{ type: 'text', text: 'run the tool' }] },
          status: 'success',
          siblingsGroupId: 0,
          createdAt: 100,
          updatedAt: 100
        },
        {
          // Assistant turn paused on a tool-approval-request — the renderer's decision arrives here.
          id: 'a1',
          parentId: 'u1',
          topicId: 'topic-1',
          role: 'assistant',
          data: {
            parts: [
              { type: 'text', text: 'let me call a tool' },
              {
                type: 'tool-fetch_url',
                toolCallId: 'call-1',
                state: 'approval-requested',
                input: { url: 'https://example.com' },
                approval: { id: APPROVAL_ID }
              }
            ]
          },
          status: 'success',
          siblingsGroupId: 1,
          modelId: ANCHOR_MODEL_ID,
          modelSnapshot: { id: 'gpt-4o-mini', name: 'GPT-4o mini', provider: 'openai' },
          createdAt: 200,
          updatedAt: 200
        }
      ])
    )
  })

  it('rejects when the anchor is not an assistant message (anchor guard)', async () => {
    await expect(
      provider.prepareDispatch(
        makeSubscriber(),
        {
          trigger: 'continue-conversation',
          topicId: 'topic-1',
          parentAnchorId: 'u1', // a user message — invalid continue anchor
          approvalDecisions: []
        },
        { hasLiveStream: false }
      )
    ).rejects.toThrow(/anchor must be an assistant message/)
  })

  it('rejects when the anchor belongs to a different topic (anchor guard)', async () => {
    await expect(
      provider.prepareDispatch(
        makeSubscriber(),
        {
          trigger: 'continue-conversation',
          topicId: 'topic-2', // anchor a1 lives on topic-1
          parentAnchorId: 'a1',
          approvalDecisions: []
        },
        { hasLiveStream: false }
      )
    ).rejects.toThrow(/anchor does not belong to topic topic-2/)
  })

  it('flips the anchor status to pending and applies the approval decision to its parts', async () => {
    await provider.prepareDispatch(
      makeSubscriber(),
      {
        trigger: 'continue-conversation',
        topicId: 'topic-1',
        parentAnchorId: 'a1',
        approvalDecisions: [{ approvalId: APPROVAL_ID, approved: true }]
      },
      { hasLiveStream: false }
    )

    const anchor = await messageService.getById('a1')
    expect(anchor.status).toBe('pending')
    const toolPart = (anchor.data.parts ?? []).find((p) => p.type === 'tool-fetch_url') as
      | { state: string; approval?: { id: string; approved?: boolean } }
      | undefined
    expect(toolPart?.state).toBe('approval-responded')
    expect(toolPart?.approval).toEqual({ id: APPROVAL_ID, approved: true })
  })

  it("reuses the anchor's model and re-anchors history on the assistant row (no new placeholder)", async () => {
    const beforeCount = (await messageService.getPathToNode('a1')).length

    const prepared = await provider.prepareDispatch(
      makeSubscriber(),
      {
        trigger: 'continue-conversation',
        topicId: 'topic-1',
        parentAnchorId: 'a1',
        approvalDecisions: [{ approvalId: APPROVAL_ID, approved: true }]
      },
      { hasLiveStream: false }
    )

    // Model reuse: the anchor's persisted modelId is what gets resolved, not the topic default.
    expect(vi.mocked(resolveModels)).toHaveBeenCalledWith([ANCHOR_MODEL_ID], MODEL_ID)

    // Single model, no sibling group, anchored back on the assistant row.
    expect(prepared.isMultiModel).toBe(false)
    expect(prepared.siblingsGroupId).toBeUndefined()
    expect(prepared.models).toHaveLength(1)
    expect(prepared.models[0].request.messageId).toBe('a1')

    // No placeholder row was created — the path to the anchor is unchanged.
    const afterCount = (await messageService.getPathToNode('a1')).length
    expect(afterCount).toBe(beforeCount)

    // History anchors on the assistant row and carries the approval-responded part.
    const history = prepared.models[0].request.messages
    expect(history?.map((m) => m.role)).toEqual(['user', 'assistant'])
    const lastAssistant = history?.[history.length - 1]
    const toolPart = lastAssistant?.parts.find((p) => p.type === 'tool-fetch_url') as { state: string } | undefined
    expect(toolPart?.state).toBe('approval-responded')
  })
})
