import type { Plugin } from "@opencode-ai/plugin"

const PROVIDER_ID = "opencode-go"
const MODEL_ID = "deepseek-v4-flash"
const MAX_ATTEMPTS = 3
const RETRY_DELAY_MS = 3000
const STATE_TTL_MS = 10 * 60 * 1000

type RetryState = {
  attempts: number
  scheduled: boolean
  lastActivity: number
}

// Keyed by `${sessionID}:${userMessageText}`: a re-send creates a new user
// message, so the text is the stable identity across retry generations.
const retryState = new Map<string, RetryState>()

const pruneStale = () => {
  const now = Date.now()
  for (const [key, state] of retryState) {
    if (now - state.lastActivity > STATE_TTL_MS && !state.scheduled) {
      retryState.delete(key)
    }
  }
}

export const RetryGoFlash: Plugin = (async ({ client }) => {
  const fetchMessage = async (sessionID: string, messageID: string) => {
    const res = await client.session.message({
      path: { id: sessionID, messageID },
    })
    if (res.error || !res.data) {
      throw new Error(res.error ? String(res.error) : "empty message response")
    }
    return res.data
  }

  const getTextParts = async (sessionID: string, userMsgID: string) => {
    const { parts } = await fetchMessage(sessionID, userMsgID)
    return parts.filter((p) => p.type === "text")
  }

  const notifyFailure = async (sessionID: string) => {
    try {
      await client.tui.showToast({
        body: {
          title: "Retry failed",
          message: `3 retries failed after stream drop (session ${sessionID}). Resume manually.`,
          variant: "error",
          duration: 10000,
        },
      })
    } catch (error) {
      console.error(`[retry-go-flash] toast failed for ${sessionID}:`, error)
    }
  }

  // Schedule one resend; on drop it re-enters via the event hook (same text
  // key, incremented attempts). If the resend itself errors out, retry it.
  const scheduleResend = (
    sessionID: string,
    userText: string,
    attempts: number,
    key: string,
    state: RetryState,
  ) => {
    if (attempts > MAX_ATTEMPTS) {
      retryState.delete(key)
      void notifyFailure(sessionID)
      return
    }
    state.scheduled = true
    state.lastActivity = Date.now()
    setTimeout(async () => {
      try {
        console.error(`[retry-go-flash] attempt ${attempts}/${MAX_ATTEMPTS} resending in ${sessionID}`)
        const res = await client.session.promptAsync({
          path: { id: sessionID },
          body: {
            parts: [{ type: "text", text: userText }],
          },
        })
        if (res.error) {
          throw new Error(String(res.error))
        }
        state.scheduled = false
        state.lastActivity = Date.now()
      } catch (error) {
        console.error(`[retry-go-flash] resend ${attempts} failed for ${sessionID}:`, error)
        scheduleResend(sessionID, userText, attempts + 1, key, state)
      }
    }, RETRY_DELAY_MS)
  }

  return {
    async event({ event }) {
      if (event.type !== "message.part.updated") return
      const part = event.properties.part
      if (part.type !== "step-finish" || part.reason !== "unknown") return
      const sessionID = part.sessionID
      const messageID = part.messageID

      pruneStale()

      // Confirm this is the flash model in build mode, not a plan/compaction abort.
      let info: Awaited<ReturnType<typeof fetchMessage>>["info"] | undefined
      try {
        const msg = await fetchMessage(sessionID, messageID)
        info = msg.info
      } catch (error) {
        console.error(`[retry-go-flash] cannot fetch message ${sessionID}/${messageID}:`, error)
        return
      }
      if (!info || info.role !== "assistant") return
      if (info.providerID !== PROVIDER_ID || !info.modelID.startsWith(MODEL_ID)) return
      if (info.summary) return
      if (info.finish !== "unknown") return
      // Real drops produce tokens; zero-token "unknown" finishes are
      // intentional aborts (permission interrupts, compaction) and must not
      // be retried.
      if (!(info.tokens?.output > 0)) return

      const userMsgID = info.parentID
      if (!userMsgID) return

      let textParts: Awaited<ReturnType<typeof getTextParts>>
      try {
        textParts = await getTextParts(sessionID, userMsgID)
      } catch (error) {
        console.error(`[retry-go-flash] cannot fetch user parts for ${sessionID}/${userMsgID}:`, error)
        return
      }
      if (textParts.length === 0) return
      const userText = (textParts[0] as { text: string }).text

      const key = `${sessionID}:${userText}`
      const existing = retryState.get(key)
      if (existing?.scheduled) return

      const attempts = (existing?.attempts ?? 0) + 1
      const state: RetryState = existing ?? { attempts: 0, scheduled: false, lastActivity: Date.now() }
      state.attempts = attempts
      retryState.set(key, state)
      scheduleResend(sessionID, userText, attempts, key, state)
    },
  }
}) satisfies Plugin