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, steps: {}, 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,
      steps: { '1:1': 10 * 1.4e-6 + 4 * 2.8e-6 },
      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,
      steps: { '1:1': 14 * 1.4e-6 + 5 * 2.8e-6 },
      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)
    expect(projected(ctx, session).steps).toEqual({ '1:1': 9 * 1.4e-6 + 1 * 2.8e-6 })
  })

  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, steps: {}, 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, steps: {}, 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, steps: { '1:1': 0 }, currency: 'USD' })
  })

  it('maps each step to its priced cost under turn:step keys', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    usageChunk(session, { inputTokens: 1_000, outputTokens: 500 }, 1, 1)
    startStep(session, 1, 2)
    usageChunk(session, { inputTokens: 200, outputTokens: 100 }, 1, 2)
    session.append('step/end', { turn: 1, step: 2 })
    expect(projected(ctx, session)).toEqual({
      totalUsd: (1_000 * 1.4e-6 + 500 * 2.8e-6) + (200 * 1.4e-6 + 100 * 2.8e-6),
      pricedSteps: 2,
      unknownModelSteps: 0,
      steps: {
        '1:1': 1_000 * 1.4e-6 + 500 * 2.8e-6,
        '1:2': 200 * 1.4e-6 + 100 * 2.8e-6,
      },
      currency: 'USD',
    })
  })

  it('drops a step entry when the final sample re-attributes it to an unknown-priced model', async () => {
    const { ctx, session } = await harness()
    recordContext(session)
    startStep(session, 1, 1)
    const source = usageChunk(session, { inputTokens: 100, outputTokens: 10 }, 1, 1)
    session.append('assistant/message', {
      turn: 1,
      step: 1,
      message: createMessage({
        role: 'assistant',
        content: [],
        source: { kind: 'model', provider: 'openrouter', model: 'brand-new/model' },
      }),
      usage: { inputTokens: 100, outputTokens: 10 },
    }, { surfaceOp: 'append', sourceEventSeqs: [source] })
    session.append('step/end', { turn: 1, step: 1 })
    // The earlier priced chunk's entry is removed with the re-attribution: the
    // step is no longer priced, so it has no entry.
    expect(projected(ctx, session)).toEqual({ totalUsd: 0, pricedSteps: 0, unknownModelSteps: 1, steps: {}, currency: 'USD' })
  })

  it('uses explicitly logged costUsd on usage directly even without pricing table entry', async () => {
    const { ctx, session } = await harness()
    session.append('request/context', { provider: 'openrouter', model: 'unlisted/custom-model' })
    startStep(session, 1, 1)
    session.append('assistant/chunk', {
      turn: 1,
      step: 1,
      chunk: {
        type: 'usage',
        usage: { inputTokens: 500, outputTokens: 200, costUsd: 0.0042 },
      },
    })
    session.append('step/end', { turn: 1, step: 1 })
    expect(projected(ctx, session)).toEqual({
      totalUsd: 0.0042,
      pricedSteps: 1,
      unknownModelSteps: 0,
      steps: { '1:1': 0.0042 },
      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('excludes the inherited prefix of a forked child, so parent + child is a sum', async () => {
    const { ctx, session: parent } = await harness()
    recordContext(parent)
    startStep(parent, 1, 1)
    const parentSource = usageChunk(parent, { inputTokens: 1_000, outputTokens: 500 }, 1, 1)
    finalUsage(parent, { inputTokens: 1_000, outputTokens: 500 }, 1, 1, [parentSource])
    const parentCost = 1_000 * 1.4e-6 + 500 * 2.8e-6
    expect(projected(ctx, parent).totalUsd).toBe(parentCost)

    // The fork: the child's log opens with a verbatim copy of the parent's.
    const seed = [...parent.events]
    const child = ctx.sessions.create(undefined, { seed, meta: { seedLength: seed.length } })
    startStep(child, 2, 1)
    const childSource = usageChunk(child, { inputTokens: 10, outputTokens: 4 }, 2, 1)
    finalUsage(child, { inputTokens: 10, outputTokens: 4 }, 2, 1, [childSource])
    const childCost = 10 * 1.4e-6 + 4 * 2.8e-6

    const childValue = projected(ctx, child)
    expect(childValue.totalUsd).toBe(childCost)
    expect(childValue.pricedSteps).toBe(1)
    expect(childValue.steps).toEqual({ '2:1': childCost })
    // The parent's own figure is untouched by the fork.
    expect(projected(ctx, parent).totalUsd).toBe(parentCost)
  })

  it('attributes a chunk-only step of a forked child from a route in the inherited prefix', async () => {
    const { ctx, session: parent } = await harness()
    // The route record is the LAST thing the parent logs, so only the
    // inherited prefix can supply the child's attribution.
    recordContext(parent)
    const seed = [...parent.events]
    const child = ctx.sessions.create(undefined, { seed, meta: { seedLength: seed.length } })
    startStep(child, 1, 1)
    usageChunk(child, { inputTokens: 100, outputTokens: 20 }, 1, 1)
    expect(projected(ctx, child)).toMatchObject({
      totalUsd: 100 * 1.4e-6 + 20 * 2.8e-6,
      pricedSteps: 1,
      unknownModelSteps: 0,
    })
  })

  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,
      steps: { '1:1': 8 * 1.4e-6 + 2 * 2.8e-6 },
      currency: 'USD',
    })
  })
})