import { Config } from "../config/config"
import z from "zod"
import { Provider } from "../provider/provider"
import { generateObject, type ModelMessage } from "ai"
import PROMPT_GENERATE from "./generate.txt"
import { SystemPrompt } from "../session/system"
import { Instance } from "../project/instance"

export namespace Agent {
  export const Info = z
    .object({
      name: z.string(),
      model: z
        .object({
          modelID: z.string(),
          providerID: z.string(),
        })
        .optional(),
      description: z.string(),
      prompt: z.string().optional(),
      tools: z.record(z.boolean()),
    })
    .openapi({
      ref: "Agent",
    })
  export type Info = z.infer<typeof Info>
  const state = Instance.state(
    async () => {
      const cfg = await Config.get()
      const result: Record<string, Info> = {
        general: {
          name: "general",
          description:
            "General-purpose agent for researching complex questions, searching for code, and executing multi-step tasks. When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries use this agent to perform the search for you.",
          tools: {
            todoread: false,
            todowrite: false,
          },
        },
      }
      for (const [key, value] of Object.entries(cfg.agent ?? {})) {
        if (value.disable) {
          delete result[key]
          continue
        }
        let item = result[key]
        if (!item)
          item = result[key] = {
            name: key,
            description: "",
            tools: {
              todowrite: false,
              todoread: false,
            },
          }
        const model = value.model ?? cfg.model
        if (model) item.model = Provider.parseModel(model)
        if (value.prompt) item.prompt = value.prompt
        if (value.tools)
          item.tools = {
            ...item.tools,
            ...value.tools,
          }
        if (value.description) item.description = value.description
      }
      return result
    },
  )

  export async function get(agent: string) {
    return state().then((x) => x[agent])
  }

  export async function list() {
    return state().then((x) => Object.values(x))
  }

  export async function generate(input: { description: string }) {
    const defaultModel = await Provider.defaultModel()
    const model = await Provider.getModel(defaultModel.providerID, defaultModel.modelID)
    const system = SystemPrompt.header(defaultModel.providerID)
    system.push(PROMPT_GENERATE)
    const existing = await list()
    const result = await generateObject({
      temperature: 0.3,
      prompt: [
        ...system.map(
          (item): ModelMessage => ({
            role: "system",
            content: item,
          }),
        ),
        {
          role: "user",
          content: `Create an agent configuration based on this request: \"${input.description}\".\n\nIMPORTANT: The following identifiers already exist and must NOT be used: ${existing.map((i) => i.name).join(", ")}\n  Return ONLY the JSON object, no other text, do not wrap in backticks`,
        },
      ],
      model: model.language,
      schema: z.object({
        identifier: z.string(),
        whenToUse: z.string(),
        systemPrompt: z.string(),
      }),
    })
    return result.object
  }
}