import { Effect } from "effect"
import type { LLM } from "@/session/llm"
import { MessageV2 } from "@/session/message-v2"
import { MessageID, PartID, type SessionID } from "@/session/schema"
import {
  type Input,
  type Chunk,
  type Output,
  type Batch,
  type Replay,
  TOOL_OUTPUT_MAX_CHARS,
  OUTPUT,
  budget,
  model,
  large,
  transcript,
  extractText,
} from "./compaction-chunks-utils"

function prompt(input: { chunk: Chunk; total: number }) {
  return [
    `Summarize conversation chunk ${input.chunk.index + 1} of ${input.total}.`,
    "Only summarize facts present in this chunk.",
    "Preserve concrete file paths, commands, errors, decisions, and unresolved tasks.",
    "Use terse Markdown bullets. Do not mention chunking or compaction.",
  ].join("\n")
}

function summarize(input: Input & { chunk: Chunk; total: number }) {
  return Effect.gen(function* () {
    const size = budget({ cfg: input.cfg, model: input.model })
    const data = (yield* large({ messages: input.chunk.messages, model: input.model, size }))
      ? [
          {
            role: "user" as const,
            content: [
              {
                type: "text" as const,
                text: [
                  "The following compacted transcript represents an oversized conversation chunk.",
                  "Summarize only facts present in the transcript.",
                  "<conversation>",
                  transcript({ messages: input.chunk.messages }),
                  "</conversation>",
                ].join("\n"),
              },
            ],
          },
        ]
      : yield* MessageV2.toModelMessagesEffect(input.chunk.messages, input.model, {
          stripMedia: true,
          toolOutputMaxChars: TOOL_OUTPUT_MAX_CHARS,
        })
    return yield* run({ ...input, data, text: prompt({ chunk: input.chunk, total: input.total }) })
  })
}

function assistant(input: { base: MessageV2.Assistant; sessionID: SessionID }) {
  return {
    ...input.base,
    id: MessageID.ascending(),
    parentID: input.base.parentID,
    sessionID: input.sessionID,
    cost: 0,
    tokens: { output: 0, input: 0, reasoning: 0, cache: { read: 0, write: 0 } },
    time: { created: Date.now() },
    finish: undefined,
    error: undefined,
  } satisfies MessageV2.Assistant
}

export function run(input: Input & { data: LLM.StreamInput["messages"]; text: string }) {
  return Effect.gen(function* () {
    const msg = yield* input.session.updateMessage(assistant({ base: input.target, sessionID: input.sessionID }))
    const mdl = model(input.model)
    const worker = yield* input.processors.create({ assistantMessage: msg, sessionID: input.sessionID, model: mdl })
    const opts = input.agent.options
    const agent = {
      ...input.agent,
      options: {
        ...opts,
        maxOutputTokens: Math.min(OUTPUT, typeof opts?.maxOutputTokens === "number" ? opts.maxOutputTokens : OUTPUT),
      },
    }
    const out = yield* Effect.gen(function* () {
      const result = yield* worker.process({
        user: input.user,
        agent,
        sessionID: input.sessionID,
        tools: {},
        system: [],
        messages: [...input.data, { role: "user", content: [{ type: "text", text: input.text }] }],
        model: mdl,
      })
      const parts = MessageV2.parts(worker.message.id)
      return { result, output: extractText(worker.message, parts) }
    }).pipe(
      Effect.ensuring(
        input.session.removeMessage({ sessionID: input.sessionID, messageID: worker.message.id }).pipe(Effect.ignore),
      ),
    )
    if (out.result !== "continue") return { result: out.result, output: undefined }
    if (!out.output) return { result: "stop" as const, output: undefined }
    return out
  })
}

function batchPrompt(input: { batch: Batch; total: number }) {
  return [
    `Summarize the middle section of conversation batch ${input.batch.index + 1} of ${input.total}.`,
    "Context before and after is provided for reference.",
    "Only summarize the content between <middle-section> tags.",
    "Preserve concrete file paths, commands, errors, decisions, and unresolved tasks.",
    "Use terse Markdown bullets. Do not mention chunking or compaction.",
  ].join("\n")
}

export function processBatch(input: Input & { b: Batch; total: number }): Effect.Effect<Output> {
  return Effect.gen(function* () {
    if (!input.b.middle.length) {
      return { result: "continue" as const, output: transcript({ messages: [...input.b.first, ...input.b.last] }) }
    }

    const size = budget({ cfg: input.cfg, model: input.model })
    const tooLarge = yield* large({ messages: input.b.middle, model: input.model, size })
    const before = input.b.first.length
      ? `<context-before>\n${transcript({ messages: input.b.first })}\n</context-before>`
      : ""
    const after = input.b.last.length
      ? `<context-after>\n${transcript({ messages: input.b.last })}\n</context-after>`
      : ""

    const data = tooLarge
      ? [
          {
            role: "user" as const,
            content: [
              {
                type: "text" as const,
                text: [
                  before,
                  `<middle-section>\n${transcript({ messages: input.b.middle })}\n</middle-section>`,
                  after,
                ].filter(Boolean).join("\n\n"),
              },
            ],
          },
        ]
      : yield* MessageV2.toModelMessagesEffect(input.b.middle, input.model, {
          stripMedia: true,
          toolOutputMaxChars: TOOL_OUTPUT_MAX_CHARS,
        })

    const promptText = [
      batchPrompt({ batch: input.b, total: input.total }),
      ...(before && !tooLarge ? ["", "Context before:", before] : []),
      ...(after && !tooLarge ? ["", "Context after:", after] : []),
    ].join("\n")

    const result = yield* run({ ...input, data, text: promptText })
    if (result.result !== "continue" || !result.output) return result

    return {
      result: "continue" as const,
      output: [
        input.b.first.length ? transcript({ messages: input.b.first }) : "",
        `<batch-summary index="${input.b.index + 1}">\n${result.output}\n</batch-summary>`,
        input.b.last.length ? transcript({ messages: input.b.last }) : "",
      ].filter(Boolean).join("\n\n"),
    }
  })
}

export function replay(input: Input & { replay: Replay }) {
  return Effect.gen(function* () {
    const chunk: Chunk = {
      index: 0,
      messages: [{ info: input.replay.info, parts: input.replay.parts }],
    }
    const size = budget({ cfg: input.cfg, model: input.model })
    if (!(yield* large({ messages: chunk.messages, model: input.model, size }))) return input.replay
    const result = yield* summarize({ ...input, chunk, total: 1 })
    if (result.result !== "continue" || !result.output) return input.replay
    return {
      info: input.replay.info,
      parts: [
        {
          id: PartID.ascending(),
          messageID: input.replay.info.id,
          sessionID: input.sessionID,
          type: "text" as const,
          synthetic: true,
          text: [
            "The original replayed request was too large to send after compaction.",
            "Use this compacted representation of that request instead:",
            result.output,
          ].join("\n\n"),
        },
      ],
    } satisfies Replay
  })
}
