import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"
import { Text } from "@earendil-works/pi-tui"

interface SpeedData {
  tps: number
  outputTokens: number
  seconds: number
}

export default function speedExtension(pi: ExtensionAPI) {
  let firstTokenAt: number | undefined
  let lastTokenAt: number | undefined

  pi.registerEntryRenderer<SpeedData>("speed", (entry, _options, theme) => {
    const { tps, outputTokens, seconds } = entry.data ?? { tps: 0, outputTokens: 0, seconds: 0 }
    return new Text(
      theme.fg("dim", `⚡ ${tps.toFixed(1)} tok/s (${outputTokens} tok in ${seconds.toFixed(1)}s)`),
    )
  })

  pi.on("message_start", (event) => {
    if (event.message.role === "assistant") {
      firstTokenAt = undefined
      lastTokenAt = undefined
    }
  })

  pi.on("message_update", (event) => {
    if (event.message.role !== "assistant") return
    const now = Date.now()
    if (firstTokenAt === undefined) firstTokenAt = now
    lastTokenAt = now
  })

  pi.on("message_end", (event) => {
    const msg = event.message
    if (msg.role !== "assistant") return
    const output = msg.usage?.output
    if (!output || firstTokenAt === undefined || lastTokenAt === undefined) return
    const seconds = (lastTokenAt - firstTokenAt) / 1000
    if (seconds <= 0) return
    pi.appendEntry<SpeedData>("speed", {
      tps: output / seconds,
      outputTokens: output,
      seconds,
    })
  })
}