import path from "path"
import fs from "fs/promises"
import { Cause, Effect, Exit } from "effect"
import { SessionID, PartID } from "@/session/schema"
import { MessageV2 } from "@/session/message-v2"
import { Session } from "@/session/session"
import { Instance } from "@/project/instance"
import type { SessionStatus } from "@/session/status"
import { Flag } from "@saeeol/core/flag/flag"
import { PlanFollowup } from "@/saeeol/plan-followup"
import type { SaeeolSession } from "@/saeeol/session"
import { Filesystem } from "@/util/filesystem"
import PROMPT_PLAN from "@/session/prompt/plan.txt"
import CODE_SWITCH from "@/session/prompt/code-switch.txt"

export function shouldAskPlanFollowup(input: { messages: MessageV2.WithParts[]; abort: AbortSignal }) {
  if (input.abort.aborted) return false
  if (!["cli", "vscode"].includes(Flag.SAEEOL_CLIENT)) return false
  const idx = input.messages.findLastIndex((m) => m.info.role === "user")
  return input.messages
    .slice(idx + 1)
    .some((msg) =>
      msg.parts.some((p) => p.type === "tool" && p.tool === "plan_exit" && p.state.status === "completed"),
    )
}

export async function askPlanFollowup(input: {
  sessionID: SessionID
  messages: MessageV2.WithParts[]
  abort: AbortSignal
}): Promise<"continue" | "break"> {
  if (!shouldAskPlanFollowup({ messages: input.messages, abort: input.abort })) return "break"
  const action = await PlanFollowup.ask({
    sessionID: input.sessionID,
    messages: input.messages,
    abort: input.abort,
  })
  return action === "continue" ? "continue" : "break"
}

export function abortPlanFollowup(sessionID: SessionID) {
  return PlanFollowup.abort(sessionID)
}

export const recoverDanglingAssistant = Effect.fn("SaeeolSessionPrompt.recoverDanglingAssistant")(function* (input: {
  sessionID: SessionID
  status: Pick<SessionStatus.Interface, "get">
  sessions: Pick<Session.Interface, "messages" | "removeMessage">
}) {
  const state = yield* input.status.get(input.sessionID)
  if (state.type !== "idle") return

  const msgs = yield* input.sessions.messages({ sessionID: input.sessionID, limit: 2 })
  const tail = msgs.at(-1)
  if (!tail || tail.info.role !== "assistant") return
  if (tail.parts.length > 0 || tail.info.finish || tail.info.error) return

  const prev = msgs.at(-2)
  if (!prev || prev.info.role !== "user") return
  if (tail.info.parentID !== prev.info.id) return

  yield* input.sessions.removeMessage({ sessionID: input.sessionID, messageID: tail.info.id })
})

export const recoverProviderFinishError = Effect.fn("SaeeolSessionPrompt.recoverProviderFinishError")(
  function* (input: {
    sessionID: SessionID
    status: Pick<SessionStatus.Interface, "get">
    sessions: Pick<Session.Interface, "messages" | "removeMessage">
  }) {
    const state = yield* input.status.get(input.sessionID)
    if (state.type !== "idle") return

    const msgs = yield* input.sessions.messages({ sessionID: input.sessionID, limit: 2 })
    const tail = msgs.at(-1)
    if (!tail || tail.info.role !== "assistant") return
    if (tail.info.finish !== "error" || tail.info.error) return
    if (!tail.parts.some((part) => part.type === "step-finish" && part.reason === "error")) return

    const prev = msgs.at(-2)
    if (!prev || prev.info.role !== "user") return
    if (tail.info.parentID !== prev.info.id) return

    yield* input.sessions.removeMessage({ sessionID: input.sessionID, messageID: tail.info.id })
  },
)

export async function ensurePlanDir(dir: string) {
  if (await Filesystem.isDir(dir)) return
  await fs.mkdir(dir, { recursive: true })
}

export async function insertPlanReminders(input: {
  agent: { name: string }
  session: Session.Info
  userMessage: MessageV2.WithParts
}) {
  if (input.agent.name !== "plan") return
  const plan = Session.plan(input.session, Instance.current)
  const exists = await Filesystem.exists(plan)
  if (!exists) await ensurePlanDir(path.dirname(plan))
  const info = exists
    ? `A plan file already exists at ${plan}. You can read it and make incremental edits using the edit tool.`
    : `No plan file exists yet. You should create your plan at ${plan} using the write tool.`
  input.userMessage.parts.push({
    id: PartID.ascending(),
    messageID: input.userMessage.info.id,
    sessionID: input.userMessage.info.sessionID,
    type: "text",
    text: PROMPT_PLAN + `\n\n## Plan File\n${info}\nThis is the ONLY file you are allowed to write to or edit.`,
    synthetic: true,
  })
}

export const CODE_SWITCH_TEXT = CODE_SWITCH

export function resolveCloseReason(input: {
  sessionID: string
  closeReasons: Map<string, SaeeolSession.CloseReason>
  exit: Exit.Exit<any, any>
}): SaeeolSession.CloseReason {
  const explicit = input.closeReasons.get(input.sessionID)
  input.closeReasons.delete(input.sessionID)
  if (explicit) return explicit
  if (Exit.isFailure(input.exit)) {
    return Cause.hasInterruptsOnly(input.exit.cause) ? "interrupted" : "error"
  }
  return "completed"
}

export const MAX_COMPACTION_ATTEMPTS = 3

export function guardCompactionAttempt(input: {
  sessionID: string
  attempts: number
  closeReasons: Map<string, SaeeolSession.CloseReason>
  message?: MessageV2.Assistant
}) {
  if (input.attempts < MAX_COMPACTION_ATTEMPTS) return { exhausted: false as const }
  const error = new MessageV2.ContextOverflowError({
    message: `Compaction exhausted: context still exceeds model limits after ${MAX_COMPACTION_ATTEMPTS} attempts`,
  }).toObject()
  input.closeReasons.set(input.sessionID, "error")
  if (input.message) {
    input.message.error ??= error
    input.message.finish ??= "error"
  }
  return { exhausted: true as const, error }
}
