import { SessionID, PartID } from "@/session/schema"
import { MessageV2 } from "@/session/message-v2"
import { Instance } from "@/project/instance"
import { Identifier } from "@/id/id"
import { environmentDetails } from "@/saeeol/editor-context"
import { StringDecoder } from "string_decoder"

export interface EnvCache {
  block?: string
  user?: string
}

export function injectEditorContext(input: {
  msgs: MessageV2.WithParts[]
  lastUser: MessageV2.User
  sessionID: SessionID
  cache: EnvCache
}) {
  if (input.cache.user !== input.lastUser.id) {
    const ctx = (() => {
      try {
        return Instance.current
      } catch {
        return undefined
      }
    })()
    input.cache.block = environmentDetails({
      ...input.lastUser.editorContext,
      ...(ctx ? { directory: ctx.directory, worktree: ctx.worktree } : {}),
    })
    input.cache.user = input.lastUser.id
  }
  if (!input.cache.block) return
  const idx = input.msgs.findLastIndex((m) => m.info.role === "user")
  if (idx === -1) return
  input.msgs[idx] = {
    ...input.msgs[idx],
    parts: [
      ...input.msgs[idx].parts,
      {
        id: PartID.make(Identifier.ascending("part")),
        sessionID: input.sessionID,
        messageID: input.msgs[idx].info.id,
        type: "text",
        text: input.cache.block,
        synthetic: true,
      } satisfies MessageV2.TextPart,
    ],
  }
}

export function createShellDecoders() {
  const stdout = new StringDecoder("utf8")
  const stderr = new StringDecoder("utf8")
  return {
    write(stream: "stdout" | "stderr", chunk: Buffer) {
      return stream === "stdout" ? stdout.write(chunk) : stderr.write(chunk)
    },
    flush() {
      return stdout.end() + stderr.end()
    },
  }
}

export function hasCompletedSummary(msgs: MessageV2.WithParts[]): boolean {
  return msgs.some((m) => m.info.role === "assistant" && m.info.summary === true && !!m.info.finish && !m.info.error)
}

export function trimBeforeLastSummary(msgs: MessageV2.WithParts[]): MessageV2.WithParts[] {
  for (let i = msgs.length - 1; i >= 0; i--) {
    const info = msgs[i].info
    if (info.role !== "assistant" || info.summary !== true || !info.finish || info.error) continue
    const parentIdx = msgs.findIndex((m) => m.info.id === info.parentID)
    if (parentIdx === -1) return msgs
    return parentIdx === 0 ? msgs : msgs.slice(parentIdx)
  }
  return msgs
}

export function stripHistoricalMedia(msgs: MessageV2.WithParts[]): MessageV2.WithParts[] {
  const cutoff = msgs.findLastIndex(
    (m) => m.info.role === "user" && m.parts.some((p) => p.type !== "text" || !p.synthetic),
  )
  if (cutoff <= 0) return msgs
  return msgs.map((msg, idx) => {
    if (idx >= cutoff) return msg
    const parts = msg.parts.map((part) => {
      if (part.type === "file" && MessageV2.isMedia(part.mime)) {
        return {
          id: part.id,
          sessionID: part.sessionID,
          messageID: part.messageID,
          type: "text" as const,
          text: `[Attached ${part.mime}: ${part.filename ?? "file"}]`,
        } satisfies MessageV2.TextPart
      }
      if (part.type === "tool" && part.state.status === "completed" && part.state.attachments?.length) {
        const kept = part.state.attachments.filter((a) => !MessageV2.isMedia(a.mime))
        if (kept.length === part.state.attachments.length) return part
        return { ...part, state: { ...part.state, attachments: kept } }
      }
      return part
    })
    return { ...msg, parts }
  })
}

export function maybeStripHistoricalMedia(msgs: MessageV2.WithParts[]): MessageV2.WithParts[] {
  return hasCompletedSummary(msgs) ? stripHistoricalMedia(msgs) : msgs
}
