import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { gzipSync } from 'node:zlib'
import { afterAll, describe, expect, it } from 'vitest'
import { newActionId, newRunId } from '../../core/ids.ts'
import type { AgentPacket } from '../../core/packet.ts'

const tmp = mkdtempSync(join(tmpdir(), 'thomas-correlate-'))
process.env.THOMAS_HOME = tmp

const { buildSessionIndex, correlateActions, msgIdFromRawRes } = await import(
  './session-correlator.ts'
)
const { appendPacket } = await import('../store/writer.ts')
const { getRawDb, closeDb } = await import('../store/db.ts')

afterAll(() => {
  closeDb()
  rmSync(tmp, { recursive: true, force: true })
})

describe('msgIdFromRawRes', () => {
  it('extracts the upstream message id from a gzipped response', () => {
    const sse =
      'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_01ABC","model":"claude-opus-4-8"}}\n\n'
    expect(msgIdFromRawRes(gzipSync(Buffer.from(sse)))).toBe('msg_01ABC')
  })
  it('returns null for absent/garbage', () => {
    expect(msgIdFromRawRes(null)).toBeNull()
    expect(msgIdFromRawRes(gzipSync(Buffer.from('no id here')))).toBeNull()
  })
})

describe('buildSessionIndex', () => {
  it('maps message.id → sessionId (+ sub-agent for sidechain turns)', () => {
    const projects = join(tmp, 'projects')
    const proj = join(projects, '-Users-x-proj')
    mkdirSync(proj, { recursive: true })
    const lines = [
      JSON.stringify({
        type: 'assistant',
        sessionId: 'sess-main',
        isSidechain: false,
        message: { id: 'msg_main1' },
      }),
      JSON.stringify({
        type: 'assistant',
        sessionId: 'sess-main',
        isSidechain: true,
        sourceToolAssistantUUID: 'sub-uuid-7',
        message: { id: 'msg_sub1' },
      }),
      JSON.stringify({ type: 'user', message: { content: 'hi' } }), // ignored
    ]
    writeFileSync(join(proj, 'sess-main.jsonl'), `${lines.join('\n')}\n`)

    const idx = buildSessionIndex(projects, Date.now() - 60_000)
    expect(idx.get('msg_main1')).toEqual({ sessionId: 'sess-main', subAgentId: null })
    expect(idx.get('msg_sub1')).toEqual({ sessionId: 'sess-main', subAgentId: 'sub-uuid-7' })
  })
})

describe('correlateActions (end to end via writer)', () => {
  it('back-fills instance_id + sub_agent_id on the action and its outcome', async () => {
    const ts = Date.now()
    const msgId = 'msg_corr_e2e'
    const rawRes = new TextEncoder().encode(
      `data: {"type":"message_start","message":{"id":"${msgId}"}}\n\n`,
    )
    const packet: AgentPacket = {
      id: newActionId(),
      runId: newRunId(),
      threadId: 'th_correlate',
      ts,
      durMs: 100,
      sourceAgent: 'claude-code',
      cost: { usd: 0.02, tokensIn: 10, tokensOut: 5 },
      rawRes,
      payload: {
        kind: 'model_call',
        protocol: 'anthropic',
        endpoint: 'https://api.anthropic.com/v1/messages',
        model: 'claude-opus-4-8',
        messages: [{ role: 'user', content: 'correlate me' }],
        stream: true,
        status: 200,
        response: [],
      },
    }
    await appendPacket(packet)
    const raw = await getRawDb()
    // Seed a matching outcome row for this action.
    raw
      .prepare(
        'INSERT INTO outcomes (action_id, run_id, thread_id, agent, ts, kind, value_usd_micro, fingerprint, verified) VALUES (?,?,?,?,?,?,?,?,1)',
      )
      .run(
        packet.id,
        packet.runId,
        packet.threadId,
        'claude-code',
        ts,
        'file_edited',
        50_000,
        'fp1',
      )

    const index = new Map([[msgId, { sessionId: 'sess-XYZ', subAgentId: 'sub-9' }]])
    const n = correlateActions(raw, index, ts - 1000)
    expect(n).toBe(1)

    const a = raw
      .prepare('SELECT instance_id, sub_agent_id FROM actions WHERE id = ?')
      .get(packet.id) as {
      instance_id: string
      sub_agent_id: string
    }
    expect(a.instance_id).toBe('sess-XYZ')
    expect(a.sub_agent_id).toBe('sub-9')
    const o = raw
      .prepare('SELECT instance_id FROM outcomes WHERE action_id = ?')
      .get(packet.id) as {
      instance_id: string
    }
    expect(o.instance_id).toBe('sess-XYZ')

    // Idempotent: re-running correlates nothing (instance_id no longer null).
    expect(correlateActions(raw, index, ts - 1000)).toBe(0)
  })
})