import type { Argv } from "yargs"
import { Instance } from "../../../project/instance"
import { Provider } from "../../../provider/provider"
import { cmd } from "../../../cli/cmd/cmd"
import { UI } from "../../../cli/ui"
import { formatTable, formatMarkdown, isTextModel } from "./roll-call-format"
import { call } from "./roll-call-call"
import type { Result } from "./roll-call-call"

export { formatTable, formatMarkdown, isTextModel } from "./roll-call-format"

export const RollCallCommand = cmd({
  command: "roll-call <filter>",
  describe: "batch-test text models matching a filter for connectivity and latency",
  builder: (yargs: Argv) => {
    return yargs
      .positional("filter", {
        type: "string",
        describe: "regex to filter models by provider/modelID (required)",
        demandOption: true,
      })
      .option("prompt", {
        type: "string",
        default: "Hello",
        describe: "Prompt to send to each model",
      })
      .option("timeout", {
        type: "number",
        default: 25000,
        describe: "Timeout for each model call in milliseconds",
      })
      .option("parallel", {
        type: "number",
        default: 5,
        describe: "Number of parallel model calls",
      })
      .option("verbose", {
        type: "boolean",
        default: false,
        describe: "Show verbose output",
      })
      .option("quiet", {
        type: "boolean",
        default: false,
        describe: "Suppress progress and decoration",
      })
      .option("output", {
        type: "string",
        choices: ["table", "json", "md"],
        default: "table",
        describe: "Output format (table, json, or md)",
      })
  },
  handler: async (args) => {
    await handle({
      prompt: args.prompt,
      timeout: args.timeout,
      filter: args.filter,
      parallel: args.parallel,
      output: args.output === "json" || args.output === "md" ? args.output : "table",
      verbose: args.verbose,
      quiet: args.quiet,
    })
  },
})

type ArgumentsCamelCase = {
  prompt: string
  timeout: number
  filter: string
  parallel: number
  output: "table" | "json" | "md"
  verbose: boolean
  quiet: boolean
  list?: typeof Provider.list
}

export async function handle(args: ArgumentsCamelCase) {
  const list = args.list ?? Provider.list

  if (args.parallel < 1) {
    UI.error("--parallel must be at least 1")
    process.exitCode = 1
    return
  }

  if (args.timeout < 1) {
    UI.error("--timeout must be at least 1")
    process.exitCode = 1
    return
  }

  if (!args.filter.trim()) {
    UI.error("filter is required and cannot be empty")
    process.exitCode = 1
    return
  }

  const json = args.output === "json"
  const structured = json || args.output === "md"
  const tty = process.stderr.isTTY ?? false
  const color = (style: string) => (tty ? style : "")

  if (!args.quiet && !structured) {
    UI.println(
      `${color(UI.Style.TEXT_INFO)}Starting roll call for models with prompt: "${args.prompt}"${color(UI.Style.TEXT_NORMAL)}`,
    )
    UI.println(
      `${color(UI.Style.TEXT_INFO)}Timeout per model: ${args.timeout}ms, Parallel calls: ${args.parallel}${color(UI.Style.TEXT_NORMAL)}`,
    )
  }

  await Instance.provide({
    directory: process.cwd(),
    async fn() {
      const providers = await list()
      const regex = (() => {
        try {
          return new RegExp(args.filter, "i")
        } catch (err) {
          UI.error(`Invalid filter regex: ${args.filter}`)
          process.exitCode = 1
          return undefined
        }
      })()
      if (!regex) return

      const models = Object.entries(providers).flatMap(([providerID, provider]) =>
        Object.entries(provider.models)
          .filter(([modelID, model]) => regex.test(`${providerID}/${modelID}`) && isTextModel(model))
          .map(([modelID, model]) => ({ providerID, modelID, model })),
      )

      if (models.length === 0) {
        if (!args.quiet && !structured)
          UI.println(`${color(UI.Style.TEXT_WARNING)}No models to test after filtering.${color(UI.Style.TEXT_NORMAL)}`)
        if (json) console.log(JSON.stringify([], null, 2))
        if (args.output === "md") console.log(formatMarkdown([]))
        if (structured) return
        process.exitCode = 1
        return
      }

      if (!args.quiet && !structured) {
        UI.println(`${color(UI.Style.TEXT_INFO)}Prompting ${models.length} models...${color(UI.Style.TEXT_NORMAL)}`)
      }

      const results: Result[] = []
      const queue = [...models]
      const active: Promise<void>[] = []

      const run = async (item: (typeof models)[0]) => {
        const name = `${item.providerID}/${item.modelID}`
        const start = Date.now()
        const result = await call(item.model, args.prompt, args.timeout, start)

        results.push({ model: name, ...result })

        if (!args.verbose || args.quiet || structured) return
        if (result.access) {
          UI.println(`${color(UI.Style.TEXT_SUCCESS)}✔${color(UI.Style.TEXT_NORMAL)} ${name} - ${result.latency}ms`)
          return
        }
        UI.println(
          `${color(UI.Style.TEXT_DANGER)}✘${color(UI.Style.TEXT_NORMAL)} ${name} - ${result.errorType}: ${result.errorMessage}`,
        )
      }

      while (queue.length > 0 || active.length > 0) {
        while (queue.length > 0 && active.length < args.parallel) {
          const item = queue.shift()
          if (!item) continue
          const promise = run(item).finally(() => {
            const index = active.indexOf(promise)
            if (index > -1) active.splice(index, 1)
          })
          active.push(promise)
        }
        if (active.length > 0) await Promise.race(active)
      }

      if (json) {
        console.log(JSON.stringify(results, null, 2))
        return
      }

      const rows = results.map((result) => [
        result.model,
        result.access ? "YES" : "NO",
        result.access ? result.snippet : result.errorMessage ? `(${result.errorMessage})` : "",
        result.latency !== null ? `${result.latency}ms` : "N/A",
      ])

      if (args.output === "md") {
        console.log(formatMarkdown(rows))
        return
      }

      const width = parseInt(process.env.COLUMNS || "", 10) || process.stdout.columns || 80
      const table = formatTable(rows, width)

      UI.println(table.header)
      UI.println(table.separator)
      table.rows.forEach((line, index) => {
        const style = results[index].access ? UI.Style.TEXT_SUCCESS : UI.Style.TEXT_DANGER
        UI.println(color(style) + line + color(UI.Style.TEXT_NORMAL))
      })

      if (args.quiet) return
      const successful = results.filter((result) => result.access).length
      const failed = results.length - successful
      UI.println("")
      UI.println(
        `${color(UI.Style.TEXT_SUCCESS)}${successful} accessible${color(UI.Style.TEXT_NORMAL)}, ${color(UI.Style.TEXT_DANGER)}${failed} failed${color(UI.Style.TEXT_NORMAL)}`,
      )
    },
  })
}
