import { Telemetry } from "@saeeol/telemetry"
import { Flag } from "@saeeol/core/flag/flag"
import { Global } from "@saeeol/core/global"
import { Instance } from "@/project/instance"
import { ProviderID, ModelID } from "@/provider/schema"
import { Question } from "@/question"
import { Session } from "@/session/session"
import { SessionID, PartID } from "@/session/schema"
import { MessageV2 } from "@/session/message-v2"
import * as Log from "@saeeol/core/util/log"
import { SaeeolSessionPromptQueue } from "@/saeeol/session/prompt-queue"
import path from "path"
import z from "zod"
import { pending, PlanFollowupRuntime, toText } from "./plan-followup-runtime"
import { startNew } from "./plan-followup-session"

export { formatTodos, generateHandover } from "./plan-followup-handover"
export { PlanFollowupRuntime } from "./plan-followup-runtime"

export namespace PlanFollowup {
  const log = Log.create({ service: "plan.followup" })

  export const PLAN_PREFIX = "Implement the following plan:"
  export const ANSWER_NEW_SESSION = "Start new session"
  export const ANSWER_CONTINUE = "Continue here"

  export function abort(sessionID: SessionID) {
    const ctl = pending.get(sessionID)
    if (!ctl) return false
    pending.delete(sessionID)
    ctl.abort()
    return true
  }

  function resolveVariant(value: string | undefined, model: { variants?: Record<string, unknown> } | undefined) {
    if (!value) return undefined
    if (!model?.variants?.[value]) return undefined
    return value
  }

  const ModelState = z
    .object({
      model: z.record(z.string(), z.object({ providerID: ProviderID.zod, modelID: ModelID.zod })).optional(),
      variant: z.record(z.string(), z.string().optional()).optional(),
    })
    .passthrough()

  async function resolveCodeModel(input: Pick<MessageV2.User, "model">) {
    const state =
      Flag.SAEEOL_CLIENT === "cli"
        ? await Bun.file(path.join(Global.Path.state, "model.json"))
            .text()
            .then((raw) => ModelState.safeParse(JSON.parse(raw)))
            .then((r) => (r.success ? r.data : undefined))
            .catch(() => undefined)
        : undefined
    const saved = state?.model?.code
    if (saved) {
      const full = await PlanFollowupRuntime.model(saved.providerID, saved.modelID).catch(() => undefined)
      if (full) {
        const key = `${saved.providerID}/${saved.modelID}`
        return { model: { ...saved, variant: resolveVariant(state?.variant?.[key], full) } }
      }
    }
    const entry = await PlanFollowupRuntime.agent("code")
    if (entry?.model) {
      const full = await PlanFollowupRuntime.model(entry.model.providerID, entry.model.modelID).catch(() => undefined)
      if (full) return { model: { ...entry.model, variant: resolveVariant(entry.variant, full) } }
    }
    return input
  }

  async function resolvePlan(input: {
    assistant?: MessageV2.WithParts
    messages: MessageV2.WithParts[]
    sessionID: SessionID
  }) {
    if (input.assistant) {
      const text = toText(input.assistant)
      if (text) return text
    }
    const lastUserIdx = input.messages.findLastIndex((m) => m.info.role === "user")
    const assistantMessages = input.messages.slice(lastUserIdx + 1).filter((m) => m.info.role === "assistant")
    const text = assistantMessages.map(toText).filter(Boolean).join("\n\n").trim()
    if (text) return text
    const session = await Session.get(SessionID.make(input.sessionID))
    const file = Bun.file(Session.plan(session, Instance.current))
    const plan = await file.text().catch(() => "")
    return plan.trim()
  }

  async function inject(input: {
    sessionID: SessionID; agent: string; model: MessageV2.User["model"]; text: string; synthetic?: boolean
  }) {
    const msg: MessageV2.User = {
      id: (await import("@/session/schema")).MessageID.ascending(),
      sessionID: input.sessionID, role: "user",
      time: { created: Date.now() }, agent: input.agent, model: input.model,
    }
    await Session.updateMessage(msg)
    await Session.updatePart({
      id: PartID.ascending(), messageID: msg.id, sessionID: input.sessionID,
      type: "text", text: input.text, synthetic: input.synthetic ?? true,
    } satisfies MessageV2.TextPart)
    return msg
  }

  function prompt(input: { sessionID: SessionID; abort: AbortSignal }) {
    const promise = PlanFollowupRuntime.question.ask({
      sessionID: input.sessionID,
      questions: [{
        question: "Ready to implement?", questionKey: "plan.followup.question",
        header: "Implement", headerKey: "plan.followup.header",
        custom: Flag.SAEEOL_CLIENT === "cli",
        options: [
          { label: ANSWER_NEW_SESSION, labelKey: "plan.followup.answer.newSession",
            description: "Implement in a fresh session with a clean context",
            descriptionKey: "plan.followup.answer.newSession.description" },
          { label: ANSWER_CONTINUE, labelKey: "plan.followup.answer.continue",
            description: "Implement the plan in this session",
            descriptionKey: "plan.followup.answer.continue.description", mode: "code" },
        ],
      }],
    })
    const listener = () =>
      PlanFollowupRuntime.question.list().then((qs) => {
        const match = qs.find((q) => q.sessionID === input.sessionID)
        if (match) PlanFollowupRuntime.question.reject(match.id)
      })
    input.abort.addEventListener("abort", listener, { once: true })
    return promise
      .catch((error) => { if (error instanceof Question.RejectedError) return undefined; throw error })
      .finally(() => { input.abort.removeEventListener("abort", listener) })
  }

  export async function ask(input: {
    sessionID: SessionID; messages: MessageV2.WithParts[]; abort: AbortSignal
  }): Promise<"continue" | "break"> {
    if (input.abort.aborted) return "break"
    const latest = input.messages.slice().reverse()
    const assistant = latest.find((msg) => msg.info.role === "assistant")
    if (!assistant) return "break"
    const plan = await resolvePlan({ assistant, messages: input.messages, sessionID: input.sessionID })
    if (!plan) return "break"
    const user = latest.find((msg) => msg.info.role === "user")?.info
    if (!user || user.role !== "user" || !user.model) return "break"

    const answers = await prompt({ sessionID: input.sessionID, abort: input.abort })
    if (!answers) { Telemetry.trackPlanFollowup(input.sessionID, "dismissed"); return "break" }
    const answer = answers[0]?.[0]?.trim()
    if (!answer) { Telemetry.trackPlanFollowup(input.sessionID, "dismissed"); return "break" }

    if (answer === ANSWER_NEW_SESSION) {
      Telemetry.trackPlanFollowup(input.sessionID, "new_session")
      await startNew({ sessionID: input.sessionID, plan, messages: input.messages,
        model: user.model, abort: input.abort, resolveCodeModel })
      return "break"
    }

    if (answer === ANSWER_CONTINUE) {
      Telemetry.trackPlanFollowup(input.sessionID, "continue")
      const code = await resolveCodeModel({ model: user.model })
      const msg = await inject({ sessionID: input.sessionID, agent: "code", model: code.model, text: "Implement the plan above." })
      SaeeolSessionPromptQueue.retarget(input.sessionID, msg.id)
      return "continue"
    }

    Telemetry.trackPlanFollowup(input.sessionID, "custom")
    const msg = await inject({ sessionID: input.sessionID, agent: "plan", model: user.model, text: answer })
    SaeeolSessionPromptQueue.retarget(input.sessionID, msg.id)
    return "continue"
  }
}
