import { Effect } from "effect"
import type { Agent } from "@/agent/agent"
import type { Config } from "@/config/config"
import type { Provider } from "@/provider/provider"
import type { SessionProcessor } from "@/session/processor"
import type { Session } from "@/session/session"
import { MessageV2 } from "@/session/message-v2"
import { usable } from "@/session/overflow"
import { type MessageID, type SessionID } from "@/session/schema"
import { Token } from "@/util/token"

export type Update = <T extends MessageV2.Part>(part: T) => Effect.Effect<T>
export type UpdateMessage = <T extends MessageV2.Info>(msg: T) => Effect.Effect<T>

export const TOOL_OUTPUT_MAX_CHARS = 2_000
export const TRANSCRIPT_MAX_CHARS = 16_000
export const RATIO = 0.6
export const CONCURRENCY = 3
export const OUTPUT = 2_048
export const DEFAULT_BATCH_SIZE = 10
export const DEFAULT_KEEP_HEAD = 2
export const DEFAULT_KEEP_TAIL = 2
export const DEFAULT_AUTO_BATCHES = 5
export const BATCH_TOKEN_TARGET = 4_000
export const BATCH_TOKEN_MAX = 8_000

export type Chunk = {
  index: number
  messages: MessageV2.WithParts[]
}

export type Output = {
  result: SessionProcessor.Result
  output: string | undefined
}

export type Deps = {
  processors: SessionProcessor.Interface
  session: Pick<Session.Interface, "updateMessage" | "updatePart" | "removeMessage">
}

export type Input = Deps & {
  user: MessageV2.User
  agent: Agent.Info
  sessionID: SessionID
  model: Provider.Model
  cfg: Config.Info
  messages: MessageV2.WithParts[]
  prompt: string
  target: MessageV2.Assistant
  updateMessage: UpdateMessage
  updatePart: Update
  auto?: boolean
  out?: { tail_start_id?: MessageID }
}

export type Replay = {
  info: MessageV2.User
  parts: MessageV2.Part[]
}

export type Batch = {
  index: number
  first: MessageV2.WithParts[]
  middle: MessageV2.WithParts[]
  last: MessageV2.WithParts[]
}

export function eligible(input: { result: SessionProcessor.Result; error: MessageV2.Assistant["error"] }) {
  if (input.result === "compact") return true
  return input.result === "stop" && input.error?.name === "ContextOverflowError"
}

export function needed(input: { cfg: Config.Info; model: Provider.Model; tokens: number }) {
  return (
    Math.ceil(input.tokens * 1.3) + model(input.model).limit.output >
    usable({ cfg: input.cfg, model: model(input.model) })
  )
}

export function budget(input: { cfg: Config.Info; model: Provider.Model }) {
  return Math.max(1_000, Math.floor(usable({ cfg: input.cfg, model: model(input.model) }) * RATIO))
}

export function model(input: Provider.Model) {
  return {
    ...input,
    limit: {
      ...input.limit,
      output: Math.min(input.limit.output, OUTPUT),
    },
  } satisfies Provider.Model
}

export function large(input: { messages: MessageV2.WithParts[]; model: Provider.Model; size: number }) {
  return estimate({ messages: input.messages, model: input.model }).pipe(Effect.map((count) => count > input.size))
}

export function estimate(input: { messages: MessageV2.WithParts[]; model: Provider.Model }) {
  return Effect.gen(function* () {
    const msgs = yield* MessageV2.toModelMessagesEffect(input.messages, input.model, {
      stripMedia: true,
      toolOutputMaxChars: TOOL_OUTPUT_MAX_CHARS,
    })
    return Token.estimate(JSON.stringify(msgs))
  })
}

export function batchSplit(messages: MessageV2.WithParts[], cfg: Config.Info): Batch[] {
  const size = cfg.compaction?.batch_size ?? DEFAULT_BATCH_SIZE
  const head = cfg.compaction?.keep_head ?? DEFAULT_KEEP_HEAD
  const tail = cfg.compaction?.keep_tail ?? DEFAULT_KEEP_TAIL
  const tokenTarget = BATCH_TOKEN_TARGET
  const tokenMax = BATCH_TOKEN_MAX
  const batches: Batch[] = []
  let i = 0
  while (i < messages.length) {
    const remaining = messages.length - i
    const fixedEnd = Math.min(i + size, messages.length)
    if (remaining <= size + head + tail) {
      const slice = messages.slice(i)
      const midStart = Math.min(head, slice.length)
      const midEnd = Math.max(midStart, slice.length - tail)
      batches.push({
        index: batches.length,
        first: slice.slice(0, midStart),
        middle: slice.slice(midStart, midEnd),
        last: slice.slice(midEnd),
      })
      break
    }
    let end = fixedEnd
    let est = 0
    for (let j = i; j < fixedEnd; j++) {
      const parts = messages[j]!.parts
      for (const p of parts) {
        if (p.type === "text") est += p.text.length
        else if (p.type === "tool" && p.state.status === "completed") est += Math.min(p.state.output?.length ?? 0, TOOL_OUTPUT_MAX_CHARS)
        else if (p.type === "reasoning") est += p.text?.length ?? 0
      }
    }
    est = Math.round(est / 4)
    if (est < tokenTarget && fixedEnd < messages.length) {
      let expandEst = est
      for (let j = fixedEnd; j < messages.length; j++) {
        const parts = messages[j]!.parts
        let add = 0
        for (const p of parts) {
          if (p.type === "text") add += p.text.length
          else if (p.type === "tool" && p.state.status === "completed") add += Math.min(p.state.output?.length ?? 0, TOOL_OUTPUT_MAX_CHARS)
        }
        expandEst += Math.round(add / 4)
        if (expandEst >= tokenTarget) { end = j + 1; break }
        if (expandEst >= tokenMax) { end = j + 1; break }
        end = j + 1
      }
    } else if (est > tokenMax) {
      for (let j = fixedEnd - 1; j > i + head; j--) {
        end = j
        break
      }
    }
    const slice = messages.slice(i, end)
    const midStart = Math.min(head, slice.length)
    const midEnd = Math.max(midStart, slice.length - tail)
    batches.push({
      index: batches.length,
      first: slice.slice(0, midStart),
      middle: slice.slice(midStart, midEnd),
      last: slice.slice(midEnd),
    })
    i = end
  }
  return batches
}

export function extractText(msg: MessageV2.Assistant, parts: MessageV2.Part[]) {
  return parts
    .filter((part): part is MessageV2.TextPart => part.type === "text" && part.messageID === msg.id)
    .map((part) => part.text.trim())
    .filter(Boolean)
    .join("\n\n")
    .trim()
}

export function clip(input: { text: string; chars: number; label: string }) {
  if (input.text.length <= input.chars) return input.text
  const cut = input.text.length - input.chars
  return `${input.text.slice(0, input.chars)}\n[${input.label} truncated for compaction: omitted ${cut} chars]`
}

export function formatPart(p: MessageV2.Part) {
  if (p.type === "text") return clip({ text: p.text, chars: TRANSCRIPT_MAX_CHARS, label: "Text" })
  if (p.type === "reasoning")
    return `[Reasoning]: ${clip({ text: p.text, chars: TRANSCRIPT_MAX_CHARS, label: "Reasoning" })}`
  if (p.type === "file") return `[File attachment]: ${p.filename ?? p.url} (${p.mime})`
  if (p.type === "agent") return `[Agent]: ${p.name}`
  if (p.type === "subtask")
    return `[Subtask ${p.agent}]: ${p.description}\n${clip({ text: p.prompt, chars: TRANSCRIPT_MAX_CHARS, label: "Subtask prompt" })}`
  if (p.type === "tool") {
    const head = `[Tool ${p.tool} ${p.state.status}]`
    if (p.state.status === "completed") {
      return [
        head,
        `input: ${clip({ text: JSON.stringify(p.state.input), chars: TOOL_OUTPUT_MAX_CHARS, label: "Tool input" })}`,
        `output: ${clip({ text: p.state.output, chars: TOOL_OUTPUT_MAX_CHARS, label: "Tool output" })}`,
      ].join("\n")
    }
    if (p.state.status === "error")
      return `${head}\n${clip({ text: p.state.error, chars: TOOL_OUTPUT_MAX_CHARS, label: "Tool error" })}`
    return `${head}\ninput: ${clip({ text: JSON.stringify(p.state.input), chars: TOOL_OUTPUT_MAX_CHARS, label: "Tool input" })}`
  }
  if (p.type === "step-finish") return `[Step finished]: ${p.reason}`
  if (p.type === "compaction") return "[Compaction requested]"
  return `[${p.type}]`
}

export function transcript(input: { messages: MessageV2.WithParts[] }) {
  return input.messages
    .map((msg, index) => {
      const body = msg.parts.map(formatPart).filter(Boolean).join("\n\n")
      return [
        `<message index="${index + 1}" role="${msg.info.role}">`,
        body || "[no content]",
        "</message>",
      ].join("\n")
    })
    .join("\n\n")
}
