import { describe, expect, test } from "bun:test"
import { Config } from "@/config/config"
import type { Provider } from "@/provider/provider"
import { batchSplit, BATCH_TOKEN_TARGET, BATCH_TOKEN_MAX } from "@/saeeol/session/compaction-chunks-utils"
import type { MessageV2 } from "@/session/message-v2"
import { MessageID, PartID, SessionID } from "@/session/schema"

function cfg(compaction?: Config.Info["compaction"]) {
  return Config.Info.zod.parse({ compaction })
}

function msg(text: string, role: "user" | "assistant" = "user"): MessageV2.WithParts {
  const sid = SessionID.make("test")
  const mid = MessageID.ascending()
  return {
    info: {
      id: mid,
      role,
      sessionID: sid,
      time: { created: Date.now() },
      parentID: role === "assistant" ? MessageID.ascending() : undefined,
    } as MessageV2.Info,
    parts: [
      {
        id: PartID.ascending(),
        messageID: mid,
        sessionID: sid,
        type: "text",
        text,
      } as MessageV2.TextPart,
    ],
  }
}

describe("compaction select budget", () => {
  test("preserveRecentBudget scales with model context size", () => {
    expect(BATCH_TOKEN_TARGET).toBeGreaterThan(0)
    expect(BATCH_TOKEN_MAX).toBeGreaterThan(BATCH_TOKEN_TARGET)
  })
})

describe("batchSplit token-aware splitting", () => {
  test("produces batches with head/tail/middle structure", () => {
    const messages = Array.from({ length: 20 }, (_, i) => msg(`message ${i}`))
    const batches = batchSplit(messages, cfg())
    expect(batches.length).toBeGreaterThan(0)
    for (const batch of batches) {
      expect(batch).toHaveProperty("first")
      expect(batch).toHaveProperty("middle")
      expect(batch).toHaveProperty("last")
      expect(batch).toHaveProperty("index")
    }
  })

  test("handles empty messages", () => {
    const batches = batchSplit([], cfg())
    expect(batches).toEqual([])
  })

  test("handles fewer messages than batch size", () => {
    const messages = [msg("a"), msg("b"), msg("c")]
    const batches = batchSplit(messages, cfg())
    expect(batches.length).toBe(1)
    expect(batches[0].first.length + batches[0].middle.length + batches[0].last.length).toBe(3)
  })

  test("respects custom batch_size config", () => {
    const longText = "x".repeat(2_000)
    const messages = Array.from({ length: 30 }, (_, i) => msg(`${longText} ${i}`))
    const batches = batchSplit(messages, cfg({ batch_size: 5 }))
    const total = batches.reduce((s, b) => s + b.first.length + b.middle.length + b.last.length, 0)
    expect(total).toBe(30)
    expect(batches.length).toBeGreaterThanOrEqual(2)
  })

  test("every message appears in exactly one batch", () => {
    const messages = Array.from({ length: 25 }, (_, i) => msg(`message ${i}`))
    const batches = batchSplit(messages, cfg())
    const total = batches.reduce((sum, b) => sum + b.first.length + b.middle.length + b.last.length, 0)
    expect(total).toBe(25)
  })

  test("expands batch when messages are short (under token target)", () => {
    const messages = Array.from({ length: 30 }, () => msg("hi"))
    const batches = batchSplit(messages, cfg({ batch_size: 5 }))
    expect(batches.length).toBeGreaterThan(0)
    const firstBatch = batches[0]
    const firstBatchTotal = firstBatch.first.length + firstBatch.middle.length + firstBatch.last.length
    expect(firstBatchTotal).toBeGreaterThan(5)
  })

  test("shrinks batch when messages are large (over token max)", () => {
    const longText = "x".repeat(40_000)
    const messages = Array.from({ length: 15 }, () => msg(longText))
    const batches = batchSplit(messages, cfg({ batch_size: 10 }))
    const firstBatch = batches[0]
    const firstBatchTotal = firstBatch.first.length + firstBatch.middle.length + firstBatch.last.length
    expect(firstBatchTotal).toBeLessThanOrEqual(10)
  })
})
