import { describe, expect, it } from 'vitest'
import { Context } from '@deepseek-ai/cordis'
import { createMessage } from '@deepseek-ai/dsh-llm'
import type { TokenUsage } from '@deepseek-ai/dsh-llm'
import SessionStore from '@deepseek-ai/dsh-session'
import type { Session } from '@deepseek-ai/dsh-session'
import SessionProjectionRegistry from '@deepseek-ai/dsh-session-projection'
import type { ModelPricing } from '@deepseek-ai/dsh-openrouter-usage/client'
import { createOpenRouterCostProjection } from '../src/projection.ts'

const PRICING = new Map<string, ModelPricing>([
  ['deepseek/deepseek-chat', {
    promptUsd: 1.4e-6,
    completionUsd: 2.8e-6,
    requestUsd: 0,
    cacheReadUsd: 1.4e-7,
    cacheWriteUsd: 1.4e-6,
  }],
  // A free model: priced, but at zero USD per bucket (still a priced step).
  ['deepseek/deepseek-chat:free', { promptUsd: 0, completionUsd: 0, requestUsd: 0 }],
  // A model with a flat per-request fee and no disclosed cache rate.
  ['expensive/request-fee', { promptUsd: 1e-5, completionUsd: 1e-5, requestUsd: 0.5 }],
])

async function harness(): Promise<{ ctx: Context; session: Session }> {
  const ctx = new Context()
  await ctx.plugin(SessionStore)
  await ctx.plugin(SessionProjectionRegistry)
  const session = ctx.sessions.create()
  ctx.sessionProjections.register(createOpenRouterCostProjection((model) => PRICING.get(model)))
  return { ctx, session }
}

function startStep(session: Session, turn: number, step: number): void {
  session.append('step/start', { turn, step })
}

/** Append a usage chunk and return its seq. */
function usageChunk(session: Session, usage: TokenUsage, turn: number, step: number): number {
  return session.append('assistant/chunk', { turn, step, chunk: { type: 'usage', usage } }).seq
}

/** Append the assistant message for a step and close it. */
function finalUsage(
  session: Session,
  usage: TokenUsage,
  turn: number,
  step: number,
  sourceSeqs: number[],
  model = 'deepseek/deepseek-chat',
): void {
  session.append('assistant/message', {
    turn,
    step,
    message: createMessage({
      role: 'assistant',
      content: [],
      source: { kind: 'model', provider: 'openrouter', model },
    }),
    usage,
  }, { surfaceOp: 'append', sourceEventSeqs: sourceSeqs })
  session.append('step/end', { turn, step })
}

function recordContext(session: Session): void {
  session.append('request/context', { provider: 'openrouter', model: 'deepseek/deepseek-chat' })
}

const projected = (ctx: Context, session: Session) => {
  const value = ctx.sessionProjections.snapshot(session).values.openRouterCost
  if (value === undefined) throw new Error('openRouterCost projection is not registered')
  return value
}

describe('openRouterCost session projection', () => {
  it('serves an all-zero view on an empty log', async () => {
    const { ctx, session } = await harness()
    expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 0, currency: 'USD' })
  })

  it('prices input/output/cache buckets at the model rates', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    const source = usageChunk(session, { inputTokens: 1_000, outputTokens: 500, cacheReadTokens: 400, cacheWriteTokens: 100 }, 1, 1)
    finalUsage(session, { inputTokens: 1_000, outputTokens: 500, cacheReadTokens: 400, cacheWriteTokens: 100 }, 1, 1, [source])
    const input = 1_000 * 1.4e-6
    const cacheRead = 400 * 1.4e-7
    const cacheWrite = 100 * 1.4e-6
    const output = 500 * 2.8e-6
    expect(projected(ctx, session).totalUsd).toBeCloseTo(input + cacheRead + cacheWrite + output, 12)
    expect(projected(ctx, session).pricedSteps).toBe(1)
  })

  it('does not double-count a usage chunk and the identical final usage', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    const source = usageChunk(session, { inputTokens: 10, outputTokens: 4 }, 1, 1)
    finalUsage(session, { inputTokens: 10, outputTokens: 4 }, 1, 1, [source])
    expect(projected(ctx, session)).toEqual({ totalUsd: 10 * 1.4e-6 + 4 * 2.8e-6, pricedSteps: 1, unknownModelSteps: 0, currency: 'USD' })
  })

  it('replaces an earlier same-step chunk sample with the final usage', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    const source = usageChunk(session, { inputTokens: 10, outputTokens: 2 }, 1, 1)
    finalUsage(session, { inputTokens: 14, outputTokens: 5 }, 1, 1, [source])
    expect(projected(ctx, session)).toEqual({
      totalUsd: 14 * 1.4e-6 + 5 * 2.8e-6,
      pricedSteps: 1,
      unknownModelSteps: 0,
      currency: 'USD',
    })
  })

  it('retains a usage chunk when no final assistant message lands (failed step)', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 9, outputTokens: 1 }, 1, 1)
    session.append('step/end', { turn: 1, step: 1 })
    expect(projected(ctx, session).totalUsd).toBeCloseTo(9 * 1.4e-6 + 1 * 2.8e-6, 12)
    expect(projected(ctx, session).pricedSteps).toBe(1)
  })

  it('counts an unknown-priced OpenRouter model as an unpriced step', async () => {
    const { ctx, session } = await harness()
    session.append('request/context', { provider: 'openrouter', model: 'brand-new/model' })
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 100, outputTokens: 1 }, 1, 1)
    session.append('step/end', { turn: 1, step: 1 })
    expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 1, currency: 'USD' })
  })

  it('ignores a step on a non-openrouter provider entirely', async () => {
    const { ctx, session } = await harness()
    session.append('request/context', { provider: 'deepseek', model: 'deepseek-chat' })
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 100, outputTokens: 1 }, 1, 1)
    session.append('step/end', { turn: 1, step: 1 })
    expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 0, currency: 'USD' })
  })

  it('prefers the assistant-message source over the last request/context record', async () => {
    const { ctx, session } = await harness()
    session.append('request/context', { provider: 'openrouter', model: 'brand-new/model' })
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 100, outputTokens: 10 }, 1, 1)
    // The message attributes the step to a priced model, overriding the last
    // request/context attribution used for the early chunk.
    const source = session.events.length - 1
    finalUsage(session, { inputTokens: 100, outputTokens: 10 }, 1, 1, [source])
    expect(projected(ctx, session).pricedSteps).toBe(1)
    expect(projected(ctx, session).unknownModelSteps).toBe(0)
  })

  it('charges the flat per-request fee once per step', async () => {
    const { ctx, session } = await harness()
    session.append('request/context', { provider: 'openrouter', model: 'expensive/request-fee' })
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 100, outputTokens: 1 }, 1, 1)
    session.append('step/end', { turn: 1, step: 1 })
    expect(projected(ctx, session).totalUsd).toBeCloseTo(100 * 1e-5 + 1 * 1e-5 + 0.5, 12)
    expect(projected(ctx, session).pricedSteps).toBe(1)
  })

  it('prices a zero-rate free model as a priced (not unknown) step', async () => {
    const { ctx, session } = await harness()
    session.append('request/context', { provider: 'openrouter', model: 'deepseek/deepseek-chat:free' })
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 100, outputTokens: 10 }, 1, 1)
    session.append('step/end', { turn: 1, step: 1 })
    expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 1, unknownModelSteps: 0, currency: 'USD' })
  })

  it('pushes no change for unrelated events', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 10, outputTokens: 1 }, 1, 1)
    const changed: string[] = []
    ctx.sessionProjections.onChanged((_session, key) => { changed.push(key) })
    session.append('todo/write', { todos: [] })
    expect(changed).not.toContain('openRouterCost')
  })

  it('restores from a JSON checkpoint', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 8, outputTokens: 2 }, 1, 1)
    session.append('step/end', { turn: 1, step: 1 })
    const checkpoint = JSON.parse(JSON.stringify(
      ctx.sessionProjections.checkpoint(session),
    )) as ReturnType<typeof ctx.sessionProjections.checkpoint>
    expect(ctx.sessionProjections.viewCheckpoint(checkpoint).openRouterCost).toEqual({
      totalUsd: 8 * 1.4e-6 + 2 * 2.8e-6,
      pricedSteps: 1,
      unknownModelSteps: 0,
      currency: 'USD',
    })
  })
})