import type { TuiPlugin, TuiPluginApi, TuiPluginModule } from "@saeeol/plugin/tui"
import { createMemo } from "solid-js"
import { formatCount, getUsage } from "@tui/routes/session/usage"
import type { AssistantMessage } from "@saeeol/sdk/v2"

const id = "internal:saeeol-sidebar-usage"

function View(props: { api: TuiPluginApi; session_id: string }) {
  const theme = () => props.api.theme.current
  const msg = createMemo(() => props.api.state.session.messages(props.session_id))
  const usage = createMemo(() => {
    const total = getUsage(msg())
    return {
      input: formatCount(total.input),
      output: formatCount(total.output),
      cached: formatCount(total.cached),
      total: total.input + total.output + total.cached,
    }
  })
  const contextInfo = createMemo(() => {
    const messages = msg()
    const last = messages.findLast((item): item is AssistantMessage => item.role === "assistant" && item.tokens.output > 0)
    if (!last) return null
    const tokens = last.tokens.input + last.tokens.output + last.tokens.reasoning + last.tokens.cache.read + last.tokens.cache.write
    const model = props.api.state.provider.find((item) => item.id === last.providerID)?.models?.[last.modelID]
    const limit = model?.limit.context ?? 0
    if (!limit) return { pct: 0, label: "", color: theme().textMuted }
    const pct = Math.round((tokens / limit) * 100)
    const color = pct >= 90 ? theme().error : pct >= 75 ? theme().warning : theme().textMuted
    const barWidth = 12
    const filled = Math.round((pct / 100) * barWidth)
    const bar = "█".repeat(filled) + "░".repeat(barWidth - filled)
    return { pct, label: `${bar} ${pct}%`, color }
  })

  return (
    <box>
      <text fg={theme().text}>
        <b>Token Usage</b>
      </text>
      <box flexDirection="row" justifyContent="space-between">
        <text fg={theme().textMuted}>Input</text>
        <text fg={theme().textMuted}>{usage().input}</text>
      </box>
      <box flexDirection="row" justifyContent="space-between">
        <text fg={theme().textMuted}>Output</text>
        <text fg={theme().textMuted}>{usage().output}</text>
      </box>
      <box flexDirection="row" justifyContent="space-between">
        <text fg={theme().textMuted}>Cached</text>
        <text fg={theme().textMuted}>{usage().cached}</text>
      </box>
      {contextInfo() && contextInfo()!.pct > 0 && (
        <box flexDirection="column">
          <text fg={theme().text}>
            <b>Context</b>
          </text>
          <text fg={contextInfo()!.color}>{contextInfo()!.label}</text>
        </box>
      )}
    </box>
  )
}

const tui: TuiPlugin = async (api) => {
  api.slots.register({
    order: 150,
    slots: {
      sidebar_content(_ctx, props) {
        return <View api={api} session_id={props.session_id} />
      },
    },
  })
}

const plugin: TuiPluginModule & { id: string } = {
  id,
  tui,
}

export default plugin
