import assert from 'node:assert/strict'
import { mkdtemp,writeFile,rm } from 'node:fs/promises'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { createContextUsageService,recoverNativeUsage } from '../src/main/context-usage-service'
import { estimateContextInput,contextInputOrigins } from '../src/main/context-usage-estimator'
import { applyAssistantEvent } from '../src/shared/message-parts'
import { contextInfoFromMessages } from '../src/shared/context-usage'
import type { ChatMessage,ContextUsageBreakdown } from '../src/shared/types'
import { Gateway } from '../src/main/gateway'

const lines=(entries:unknown[])=>entries.map(entry=>JSON.stringify(entry)).join('\n')
const user={id:'native-user',type:'message',message:{role:'user',content:[{type:'text',text:'hello'}]}}
const first={id:'native-a1',parentId:'native-user',type:'message',message:{role:'assistant',model:'fixture',usage:{input:50,output:10},stopReason:'toolUse',timestamp:100}}
const last={id:'native-a2',parentId:'native-a1',type:'message',message:{role:'assistant',model:'fixture',usage:{input:4009,cacheRead:1152,output:125,totalTokens:5286},stopReason:'stop',timestamp:200}}
const history=lines([user,first,last])
assert.equal(recoverNativeUsage(history,'native-user')?.sample.usage.inputTokens,5161,'cache input belongs to input occupancy, not output')
assert.equal(recoverNativeUsage(history,'native-user')?.sample.usage.totalTokens,5286)
assert.equal(recoverNativeUsage(history,'native-user')?.sample.nativeEntryId,'native-a2','use the latest request, never the sum of tool rounds')
assert.equal(recoverNativeUsage(lines([user,first,last,{...last,id:'sibling',parentId:'native-user'}]),'native-user'),undefined,'sibling histories cannot be guessed')
assert.equal(recoverNativeUsage(lines([user,first,{...last,message:{...last.message,stopReason:'error'}}]),'native-user'),undefined,'a failed final request does not inherit earlier tool usage')

const directory=await mkdtemp(join(tmpdir(),'ycode-context-'))
try {
 const path=join(directory,'native.jsonl')
 await writeFile(path,history)
 let messages:ChatMessage[]=[{id:'user',role:'user',content:'hello',createdAt:1,requestId:'run',nativeEntryId:'native-user',nativeSessionPath:path},
   {id:'assistant',role:'assistant',content:'done',createdAt:2,requestId:'run',modelId:'model',status:'done'}]
 let persisted=0,running=false
 let live:(()=>Promise<unknown>)|undefined
 const service=createContextUsageService({messages:async()=>messages,modelId:async()=> 'model',nativePath:()=>path,running:()=>running,
  liveState:()=>live?.(),persistIfCurrent:async(_id,assistantId,requestId,sample)=>{
   const current=messages.findLast(message=>message.role==='assistant')
   if(!running&&current?.id===assistantId&&current.requestId===requestId){current.contextUsage=sample;persisted++}
  }})
 const restored=await service.get('session')
 assert.equal(restored.sample?.usage.inputTokens,5161)
 assert.equal(restored.sample?.source,'native-history')
 assert.equal(restored.sample?.contextWindow,undefined,'do not invent a historical denominator from current settings')
 assert.equal(persisted,1)
 await service.get('session');assert.equal(persisted,1,'reading a cached sample does not rewrite history repeatedly')
 messages[1].contextUsage=undefined
 live=async()=>({success:true,data:{sessionFile:path,model:{id:'fixture'},contextUsage:{tokens:5161,contextWindow:64000}}})
 assert.equal((await service.get('session')).sample?.contextWindow,64000,'matching live native state can supply its window')
 messages[1].contextUsage=undefined
 messages[0].nativeSessionPath=join(directory,'outside.jsonl')
 assert.equal((await service.get('session')).sample,undefined,'a message cannot redirect the reader to an arbitrary file')
 messages[0].nativeSessionPath=path
 running=true
 assert.equal((await service.get('session')).sample,undefined,'querying a running turn must not recover a previous round')
 running=false
 let release!:(value:unknown)=>void
 live=()=>new Promise(resolve=>{release=resolve})
 const stale=service.get('session')
 while(!release)await new Promise(resolve=>setTimeout(resolve,1))
 messages=[...messages,{id:'new-user',role:'user',content:'next',createdAt:3},{id:'new-assistant',role:'assistant',content:'',createdAt:4,status:'streaming',requestId:'next'}]
 release({data:{sessionFile:path,model:{id:'fixture'},contextUsage:{tokens:5161,contextWindow:64000}}})
 assert.equal((await stale).sample,undefined,'a delayed query cannot overwrite a newer turn')
 assert.equal(persisted,2)
 assert.equal(contextInfoFromMessages(messages).sample,undefined,'no traversal past the latest assistant')
}finally{await rm(directory,{recursive:true,force:true})}

const skill='<skill id="skill:test">skill instructions</skill>'
const injected=`Project rules\n\n${skill}`
const input={systemPrompt:'Base system instructions',origins:contextInputOrigins([injected]),
 tools:[{type:'function',function:{name:'unused_tool',description:'This declaration is present without a call',parameters:{type:'object'}}},
  {name:'mcp_fixture_lookup',description:'MCP lookup',parameters:{type:'object'}}],
 messages:[{role:'user',content:`Question\n\n${injected}`},
  {role:'assistant',tool_calls:[{id:'call',function:{name:'mcp_fixture_lookup',arguments:'{}'}}]},
  {role:'tool',tool_call_id:'call',content:'returned MCP contents'}]}
const counted=estimateContextInput(input)
assert.ok(counted.counts.system>0&&counted.counts.skills>0&&counted.counts.tools>0&&counted.counts.conversation>0&&counted.counts.mcp>0)
assert.equal(Object.values(counted.counts).reduce((a,b)=>a+b,0),counted.inputTokens)
assert.equal(estimateContextInput({messages:[{role:'user',content:skill}]}).counts.skills,0,'ordinary user text is not relabeled by a tag pattern')
assert.equal(estimateContextInput({messages:[{role:'user',content:[{type:'image_url',image_url:{url:'data:image/png;base64,secret'}}]}]}).partial,true,'images are partial, not base64 tokens')
const breakdown:ContextUsageBreakdown={...counted,requestId:'run',sequence:1,sampledAt:1,method:'utf8-estimate',scope:'provider-request'}
let assistant:ChatMessage={id:'a',role:'assistant',content:'',createdAt:0,requestId:'run'}
assistant=applyAssistantEvent(assistant,{runId:'run',type:'context_breakdown',breakdown})
assistant=applyAssistantEvent(assistant,{runId:'run',type:'usage',usage:{inputTokens:9999,outputTokens:1,totalTokens:10000}})
assert.equal(assistant.contextUsage?.usage.inputTokens,9999)
assert.equal(assistant.contextBreakdown?.inputTokens,counted.inputTokens,'classification is not rescaled to provider usage')
assistant=applyAssistantEvent(assistant,{runId:'run',type:'context_breakdown',breakdown})
assert.equal(assistant.contextUsage?.usage.inputTokens,9999,'a repeated old observer frame cannot clear newer usage')
assistant=applyAssistantEvent(assistant,{runId:'run',type:'context_breakdown',breakdown:{...breakdown,sequence:2}})
assert.equal(assistant.contextUsage,undefined,'a new prepared request invalidates the previous provider total')
assert.equal(contextInfoFromMessages([{...assistant,contextBreakdown:breakdown},{id:'pending',role:'user',content:'next',createdAt:5}]).status,'previous',
 'an unmatched latest user message marks any retained statistic as previous')
assert.equal(contextInfoFromMessages([{...assistant,requestId:'different',contextBreakdown:breakdown}]).breakdown,undefined,
 'a stored classification for a different request is not reused')

// The fallback observes the final body for both supported provider protocols.
const originalFetch=globalThis.fetch
try {
 for(const responses of [false,true]) {
  let observed:Record<string,unknown>|undefined,transmitted:Record<string,unknown>|undefined
  globalThis.fetch=(async(_url,options)=>{
   transmitted=JSON.parse(String(options?.body))
   const payload=responses
    ? 'data: '+JSON.stringify({type:'response.completed',response:{id:'reply',status:'completed',output:[],usage:{input_tokens:10,output_tokens:1,total_tokens:11}}})+'\n\n'
    : 'data: '+JSON.stringify({choices:[{index:0,delta:{content:'done'},finish_reason:'stop'}],usage:{prompt_tokens:10,completion_tokens:1,total_tokens:11}})+'\n\ndata: [DONE]\n\n'
   return new Response(payload,{status:200,headers:{'content-type':'text/event-stream'}})
  }) as typeof fetch
  const gateway=new Gateway({baseUrl:responses?'http://127.0.0.1:1/v1/responses':'http://127.0.0.1:1/v1',apiKey:'',authMode:'none',model:'fixture',temperature:0,maxToolRounds:1},body=>{observed=structuredClone(body)})
  for await(const _delta of gateway.stream({messages:[{role:'system',content:'system'},{role:'user',content:'hello'}],tools:[{type:'function',function:{name:'unused',description:'definition',parameters:{type:'object'}}}]})) {/* fully consume */}
  assert.deepEqual(observed,transmitted,'observer sees the exact transmitted body without changing it')
  const body=observed as any
  assert.ok(estimateContextInput({messages:body.messages??body.input,systemPrompt:body.instructions,tools:body.tools}).counts.tools>0)
 }
}finally{globalThis.fetch=originalFetch}
console.log('context-usage-regression: all checks passed')