import assert from 'node:assert/strict'
import { createContextUsageService, contextSnapshotMatches, type ContextUsageDependencies } from '../src/main/context-usage-service'
import { contextBreakdownMethodLabel } from '../src/shared/context-usage'
import type { ChatMessage, ContextUsageBreakdown } from '../src/shared/types'

const sample = {
  usage: { inputTokens: 1200, outputTokens: 100, totalTokens: 1300 },
  contextWindow: 10_000,
  source: 'upstream' as const,
  sampledAt: 100,
  requestId: 'request-a',
  modelId: 'model-a'
}
const assistant = (overrides: Partial<ChatMessage> = {}): ChatMessage => ({
  id: 'assistant-a',
  role: 'assistant',
  content: 'answer',
  modelId: 'model-a',
  requestId: 'request-a',
  contextUsage: sample,
  createdAt: 1,
  status: 'done',
  ...overrides
})
const breakdown: ContextUsageBreakdown = {
  requestId: 'request-a',
  sequence: 3,
  sampledAt: 101,
  modelId: 'model-a',
  contextWindow: 10_000,
  method: 'native-summary',
  scope: 'provider-request',
  counts: { system: 100, tools: 100, conversation: 900, mcp: 50, skills: 50, other: 0 },
  inputTokens: 1200,
  partial: false
}

function service(overrides: Partial<ContextUsageDependencies> = {}) {
  const messages = [assistant()]
  const dependencies: ContextUsageDependencies = {
    messages: async () => messages,
    modelId: async () => 'model-a',
    running: () => false,
    summary: async () => undefined,
    persistIfCurrent: async (_id, guard, summary) => {
      const current = await dependencies.messages(_id)
      if(contextSnapshotMatches(current, guard, summary))current.findLast(message=>message.role==='assistant')!.contextBreakdown=summary
    },
    ...overrides
  }
  return createContextUsageService(dependencies)
}

assert.deepEqual(await service({ messages: async () => [] }).get('empty'),
  { status: 'unavailable', reason: '发送消息后可查看用量' })

const latest = await service().get('latest')
assert.equal(latest.status, 'latest')
assert.equal(latest.sample?.usage.inputTokens, 1200)

const previous = await service({ modelId: async () => 'model-b' }).get('previous')
assert.equal(previous.status, 'previous')
assert.equal(previous.reason, '来自切换模型前的最近请求')

const running = await service({ modelId: async () => 'model-b', running: () => true }).get('running')
assert.equal(running.status, 'latest', 'the active run must not be relabelled when the next-send model changes')

const summary = await service({ summary: async () => breakdown }).get('summary')
assert.equal(summary.breakdown?.sequence, 3)
assert.equal(summary.status, 'latest', 'a matching native summary augments the current request without replacing its upstream total')

const mismatched = await service({ summary: async () => ({ ...breakdown, requestId: 'older-request' }) }).get('mismatch')
assert.equal(mismatched.breakdown, undefined, 'a late summary for another request must not attach to the latest assistant message')

const summaryFailure = await service({ summary: async () => { throw new Error('query failed') } }).get('failure')
assert.equal(summaryFailure.sample?.usage.inputTokens, 1200, 'summary failures retain the real streamed usage sample')

console.log('context-usage-regression: latest, previous, active-run, summary identity and failure cases passed')

const cachedMessages = [assistant()]
const originalUsage = JSON.stringify(cachedMessages[0].contextUsage)
await service({messages:async()=>cachedMessages,summary:async()=>breakdown}).get('save')
assert.deepEqual(cachedMessages[0].contextBreakdown,breakdown)
assert.equal(JSON.stringify(cachedMessages[0].contextUsage),originalUsage,'saving categories never overwrites actual upstream usage')
const coldMessages=JSON.parse(JSON.stringify(cachedMessages)) as ChatMessage[]
for(const summary of [async()=>undefined,async()=>{throw new Error('kernel shut down')}]) {
  const offline=await service({messages:async()=>coldMessages,summary}).get('offline')
  assert.deepEqual(offline.breakdown,breakdown,'persisted categories survive cold/offline/no-instance queries')
  assert.equal(offline.sample?.contextWindow,10_000)
}
const oldModel=await service({messages:async()=>coldMessages,modelId:async()=>'model-b'}).get('switch')
assert.equal(oldModel.status,'previous')
assert.equal(oldModel.sample?.contextWindow,10_000,'never combine old sample with the new selected model window')
assert.equal((await service({summary:async()=>({...breakdown,modelId:'wrong'})}).get('wrong-model')).breakdown,undefined)
for(const change of ['assistant','rewind','pending-user'] as const) {
  let current=[assistant()]
  const result=await service({messages:async()=>current,summary:async()=>{
    if(change==='assistant')current.push(assistant({id:'new',requestId:'new-request',contextUsage:undefined}))
    if(change==='rewind')current=[{id:'different-user',role:'user',content:'fork',createdAt:0},assistant()]
    if(change==='pending-user')current.push({id:'pending',role:'user',content:'next',createdAt:2})
    return breakdown
  }}).get(change)
  assert.equal(result.breakdown,undefined,`${change} while summary awaits cannot pollute the active branch`)
}
const pending=[assistant({contextBreakdown:breakdown}),{id:'pending',role:'user',content:'next',createdAt:2} as ChatMessage]
assert.equal((await service({messages:async()=>pending,summary:async()=>breakdown}).get('pending')).status,'previous','refresh must not erase pending-user previous status')
const categoriesOnly=await service({messages:async()=>[assistant({contextUsage:undefined,contextBreakdown:breakdown})]}).get('categories-only')
assert.equal(categoriesOnly.status,'latest','categories without upstream usage remain available, not fake zero usage')
assert.equal(categoriesOnly.sample,undefined)
assert.equal(contextBreakdownMethodLabel('native-summary'),'内核 summary 本地估计')
assert.equal(contextBreakdownMethodLabel('utf8-estimate'),'UTF-8 字节估计')
console.log('context-usage persistence: offline, shutdown, identity races, usage preservation, category provenance passed')