register.tsx

import type { EngineInterface, Register, RenderElement, RenderSurface, Timer } from 'claude-code'

import type { Auto, OffReason, SavedSession, Usage } from '../types'
import {
  isWarningVisible,
  nextStep,
  offToast,
  readSettings,
  remainingMs,
  statusLayout,
  toUsage,
  warningLayout,
  warningToast,
} from './logic'
import type { Control, Part, Segment, Settings, Snapshot, Tone } from './logic'

// Vorgaben der Einstellungen. Ändern lassen sie sich in /config (userConfig des
// Mods); diese Werte gelten, solange dort nichts eingestellt ist.
const DEFAULTS = {
  cacheMinutes: 60, // Cache-Dauer: 60 bei Claude-Abos, 5 per API
  pingLeadMinutes: 3, // Vorlauf des automatischen Pings vor Ablauf
  maxHours: 4, // Höchstdauer des Warmhaltens ab dem Einschalten
  limitPercent: 90, // ab diesem Stand des 5-Stunden-Limits schaltet es sich aus
  warnLeadMinutes: 5, // Vorlauf der Warnung vor Ablauf
}

const TICK_MS = 15_000
const PING_PROMPT =
  'Cache-Ping von keepcachewarm: Bitte bestätige nur mit dem Wort „ok“ und tu sonst nichts.'
const SOUND = 'sounds/warnung.wav'
const STORE_PREFIX = 'session:'
const STORE_KEEP = 30
const AUTO_OFF: Auto = { isOn: false, since: null, offReason: null }

const NOW = { plugin: 'keepcachewarm', key: 'now' } as const
const SESSION_ID = { plugin: 'keepcachewarm', key: 'sessionId' } as const
const IS_WORKING = { plugin: 'keepcachewarm', key: 'isWorking' } as const
const LAST_ANSWER_AT = { plugin: 'keepcachewarm', key: 'lastAnswerAt' } as const
const PINGED_FOR = { plugin: 'keepcachewarm', key: 'pingedFor' } as const
const WARNED_FOR = { plugin: 'keepcachewarm', key: 'warnedFor' } as const
const IS_PING_WAITING = { plugin: 'keepcachewarm', key: 'isPingWaiting' } as const
const AUTO = { plugin: 'keepcachewarm', key: 'auto' } as const
const USAGE = { plugin: 'keepcachewarm', key: 'usage' } as const

type Engine = EngineInterface

/** Was der Mod über die laufende Session weiß; `$.state` spiegelt es fürs Zeichnen. */
type Model = {
  sessionId: string | null
  isWorking: boolean
  lastAnswerAt: number | null
  pingedFor: number | null
  warnedFor: number | null
  isPingWaiting: boolean
  auto: Auto
  usage: Usage | null
}

const EMPTY: Model = {
  sessionId: null,
  isWorking: false,
  lastAnswerAt: null,
  pingedFor: null,
  warnedFor: null,
  isPingWaiting: false,
  auto: AUTO_OFF,
  usage: null,
}

// Was das Modul selbst hält; ein Neuladen beginnt hier neu, `$.state` behält den Rest.
let cfg: Settings = readSettings({}, DEFAULTS)
let model: Model | null = null
let loading: Promise<Model> | null = null
let evaluating: Promise<void> | null = null
let isEvaluateAgain = false
let clock: Timer | null = null
// Ob im Eingabefeld Text steht, wie die Hinweiszeile unter dem Feld es je Oberfläche meldet.
const drafts = new Map<RenderSurface, boolean>()

// -------------------------------------------------------------------------
// Zustand

async function save($: Engine, patch: Partial<Model>): Promise<Model> {
  const next: Model = { ...(model ?? EMPTY), ...patch }
  model = next
  if ('sessionId' in patch) await $.state.set(SESSION_ID, next.sessionId)
  if ('isWorking' in patch) await $.state.set(IS_WORKING, next.isWorking)
  if ('lastAnswerAt' in patch) await $.state.set(LAST_ANSWER_AT, next.lastAnswerAt)
  if ('pingedFor' in patch) await $.state.set(PINGED_FOR, next.pingedFor)
  if ('warnedFor' in patch) await $.state.set(WARNED_FOR, next.warnedFor)
  if ('isPingWaiting' in patch) await $.state.set(IS_PING_WAITING, next.isPingWaiting)
  if ('auto' in patch) await $.state.set(AUTO, next.auto)
  if ('usage' in patch) await $.state.set(USAGE, next.usage)
  return next
}

/** Nach einem Neuladen des Mods aus `$.state`, sonst aus dem Speicher der Session. */
function ensureModel($: Engine): Promise<Model> {
  if (model !== null) return Promise.resolve(model)
  loading ??= (async () => {
    const id = await $.session.id()
    const { value: heldId = null } = await $.state.get(SESSION_ID)
    if (heldId === id) {
      model = {
        sessionId: id,
        isWorking: (await $.state.get(IS_WORKING)).value ?? false,
        lastAnswerAt: (await $.state.get(LAST_ANSWER_AT)).value ?? null,
        pingedFor: (await $.state.get(PINGED_FOR)).value ?? null,
        warnedFor: (await $.state.get(WARNED_FOR)).value ?? null,
        isPingWaiting: (await $.state.get(IS_PING_WAITING)).value ?? false,
        auto: (await $.state.get(AUTO)).value ?? AUTO_OFF,
        usage: (await $.state.get(USAGE)).value ?? null,
      }
      return model
    }
    return loadSession($, id)
  })().finally(() => {
    loading = null
  })
  return loading
}

async function loadSession($: Engine, id: string): Promise<Model> {
  const saved = asSaved(await $.store.get(STORE_PREFIX + id))
  const now = await $.clock.now()
  let auto = saved?.auto ?? AUTO_OFF
  // Das Warmhalten übersteht ein Fortsetzen nur, solange die Höchstdauer nicht um ist.
  if (auto.isOn && (auto.since === null || now - auto.since >= cfg.maxOnMs)) {
    auto = { isOn: false, since: null, offReason: 'max-duration' }
  }
  await save($, {
    sessionId: id,
    isWorking: model?.isWorking ?? false,
    lastAnswerAt: saved?.lastAnswerAt ?? null,
    pingedFor: saved?.pingedFor ?? null,
    warnedFor: saved?.warnedFor ?? null,
    isPingWaiting: false,
    auto,
  })
  await refreshUsage($)
  return model ?? EMPTY
}

async function persist($: Engine): Promise<void> {
  if (model?.sessionId == null) return
  const saved: SavedSession = {
    lastAnswerAt: model.lastAnswerAt,
    pingedFor: model.pingedFor,
    warnedFor: model.warnedFor,
    auto: model.auto,
    savedAt: await $.clock.now(),
  }
  await $.store.set(STORE_PREFIX + model.sessionId, saved)

  // Nur die jüngsten Sessions behalten, damit der Speicher nicht wächst.
  const keys = (await $.store.keys()).filter(key => key.startsWith(STORE_PREFIX))
  if (keys.length <= STORE_KEEP) return
  const dated = await Promise.all(
    keys.map(async key => ({ key, savedAt: asSaved(await $.store.get(key))?.savedAt ?? 0 })),
  )
  dated.sort((a, b) => a.savedAt - b.savedAt)
  for (const { key } of dated.slice(0, dated.length - STORE_KEEP)) {
    await $.store.delete(key)
  }
}

async function refreshUsage($: Engine): Promise<void> {
  try {
    await save($, { usage: toUsage(await $.session.usage()) })
  } catch {
    // Ohne frische Werte bleiben die alten stehen.
  }
}

/** Neue Session-ID (z. B. nach /resume): deren gemerkten Stand laden. */
async function syncSession($: Engine): Promise<Model> {
  const current = await ensureModel($)
  const id = await $.session.id()
  return current.sessionId === id ? current : loadSession($, id)
}

function snapshotOf(m: Model, now: number): Snapshot {
  return {
    now,
    isWorking: m.isWorking,
    lastAnswerAt: m.lastAnswerAt,
    pingedFor: m.pingedFor,
    warnedFor: m.warnedFor,
    auto: m.auto,
    usage: m.usage,
  }
}

// -------------------------------------------------------------------------
// Warnen, pingen, selbst abschalten

/** Im Zweifel steht Text im Feld: dann wird nicht gepingt. */
async function hasDraft($: Engine): Promise<boolean> {
  try {
    if ((await $.prompt.read()).text.trim() !== '') return true
    const surfaces = await $.session.surfaces()
    return surfaces.some(surface => drafts.get(surface) === true)
  } catch {
    return true
  }
}

async function sendPing($: Engine, window: number | null): Promise<void> {
  await save($, { pingedFor: window, isPingWaiting: false })
  await persist($)
  void $.prompt.submit({ text: PING_PROMPT }).catch(() => undefined)
}

async function switchOff($: Engine, reason: OffReason, s: Snapshot): Promise<void> {
  const m = await save($, {
    auto: { isOn: false, since: null, offReason: reason },
    isPingWaiting: false,
  })
  await persist($)
  $.ui.toast(offToast(reason, snapshotOf(m, s.now), cfg), { timeoutMs: 8000 })
}

async function warn($: Engine, s: Snapshot): Promise<void> {
  await save($, { warnedFor: s.lastAnswerAt })
  await persist($)
  $.ui.toast(warningToast(remainingMs(s, cfg) ?? 0), { timeoutMs: 10_000 })
  void $.audio.play({ asset: SOUND }).catch(() => undefined)
}

async function evaluateOnce($: Engine): Promise<void> {
  let m = await ensureModel($)
  const now = await $.clock.now()
  // Nach einer Selbstabschaltung noch einmal fragen: dann kann die Warnung fällig sein.
  for (let round = 0; round < 3; round += 1) {
    const s = snapshotOf(m, now)
    const step = nextStep(s, cfg)
    if (step.kind === 'switch-off') {
      await switchOff($, step.reason, s)
      m = model ?? m
      continue
    }
    if (step.kind === 'warn') {
      await warn($, s)
    } else if (step.kind === 'ping') {
      if (await hasDraft($)) {
        if (!m.isPingWaiting) await save($, { isPingWaiting: true })
      } else {
        await sendPing($, s.lastAnswerAt)
      }
    } else if (m.isPingWaiting) {
      await save($, { isPingWaiting: false })
    }
    return
  }
}

/** Eine Prüfung zur Zeit; was währenddessen angestoßen wird, läuft gleich danach. */
function evaluate($: Engine): Promise<void> {
  if (evaluating !== null) {
    isEvaluateAgain = true
    return evaluating
  }
  evaluating = (async () => {
    do {
      isEvaluateAgain = false
      await evaluateOnce($)
    } while (isEvaluateAgain)
  })().finally(() => {
    evaluating = null
  })
  return evaluating
}

async function tick($: Engine): Promise<void> {
  await syncSession($)
  await $.state.set(NOW, await $.clock.now())
  await evaluate($)
}

// -------------------------------------------------------------------------
// Knöpfe

async function toggleAuto($: Engine): Promise<void> {
  const m = await ensureModel($)
  const auto: Auto = m.auto.isOn
    ? AUTO_OFF
    : { isOn: true, since: await $.clock.now(), offReason: null }
  await save($, { auto, isPingWaiting: false })
  await persist($)
  await evaluate($)
}

async function pingOnce($: Engine): Promise<void> {
  const m = await ensureModel($)
  if (m.lastAnswerAt === null || m.isWorking) return
  await sendPing($, m.lastAnswerAt)
}

function compact($: Engine): void {
  void $.command.run({ command: 'compact', args: '' }).catch(() => undefined)
}

export const register: Register = (on, options) => {
  cfg = readSettings(options, DEFAULTS)

  // -------------------------------------------------------------------------
  // Ereignisse

  on('session.start', async ($, e, next) => {
    const result = await next(e)
    await ensureModel($)
    await refreshUsage($)
    await tick($)
    clock?.cancel()
    clock = $.clock.every(TICK_MS, () => void tick($).catch(() => undefined))
    return result
  })

  on('session.end', async ($, e, next) => {
    const result = await next(e)
    // Nach /clear läuft der Prozess unter neuer Session-ID weiter, ohne Cache.
    if (e.reason === 'clear') {
      await save($, { ...EMPTY, usage: model?.usage ?? null })
    }
    return result
  })

  on('turn.start', async ($, e, next) => {
    await ensureModel($)
    await save($, { isWorking: true, isPingWaiting: false })
    return next(e)
  })

  on('turn.complete', async ($, e, next) => {
    const result = await next(e)
    // Antworten von Subagenten öffnen kein Fenster für die Hauptsession.
    if (e.agentId !== undefined) return result

    await ensureModel($)
    const now = await $.clock.now()
    const isAnswered = e.reason === 'answer' || e.usage !== undefined
    await save($, isAnswered ? { isWorking: false, lastAnswerAt: now } : { isWorking: false })
    await refreshUsage($)
    await persist($)
    await $.state.set(NOW, now)
    await evaluate($)
    return result
  })

  // Nach dem Komprimieren ist der Kontext kleiner: Werte neu einlesen. Scheitert
  // das, läuft das Komprimieren trotzdem durch.
  on('session.compact', async ($, e, next) => {
    const result = await next(e)
    await ensureModel($)
    await refreshUsage($)
    return result
  }).catch(($, e, next) => next(e))

  on('ui.render', { component: 'PromptHint' }, ($, e, next) => {
    drafts.set(e.surface, e.props.isDraft)
    return next(e)
  })

  // -------------------------------------------------------------------------
  // Zeichnen: eine Zeile über dem Eingabefeld, darüber bei Bedarf die Warnung

  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
    // Eine Umfrage hat Vorrang: dann nichts zusätzlich zeichnen.
    if (e.props.hasSurvey) return next(e)
    // Zeilen anderer Mods bleiben stehen, unsere kommt darunter.
    const others = await next(e)

    const held = (await $.state.get(NOW)).value ?? 0
    const now = held > 0 ? held : await $.clock.now()
    const s: Snapshot = {
      now,
      isWorking: e.props.isWorking || ((await $.state.get(IS_WORKING)).value ?? false),
      lastAnswerAt: (await $.state.get(LAST_ANSWER_AT)).value ?? null,
      pingedFor: (await $.state.get(PINGED_FOR)).value ?? null,
      warnedFor: (await $.state.get(WARNED_FOR)).value ?? null,
      auto: (await $.state.get(AUTO)).value ?? AUTO_OFF,
      usage: (await $.state.get(USAGE)).value ?? null,
    }
    const isPingWaiting = (await $.state.get(IS_PING_WAITING)).value ?? false

    const { Box, Button, Text } = $.ui.resolve(e)
    const isWarning = isWarningVisible(s, cfg) && s.pingedFor !== s.lastAnswerAt
    const layout = statusLayout(s, cfg, isPingWaiting, isWarning, e.props.bodyColumns)

    const text = (part: Part) => (
      <Text {...toneProps(part.tone)} wrap="truncate">
        {part.text}
      </Text>
    )
    const segment = (one: Segment, index: number) => {
      const box = (
        <Box key={one.key} flexShrink={0}>
          {one.parts.map(text)}
        </Box>
      )
      return index === 0 ? [box] : [<Text dimColor>{' │ '}</Text>, box]
    }
    const control = (one: Control) => (
      <Box key={one.key === 'auto' ? 'schalter' : 'komprimieren'} marginLeft={2} flexShrink={0}>
        <Button
          key={one.key}
          label={one.label}
          hotkey={one.key === 'auto' ? 'w' : 'k'}
          {...(one.isPrimary ? { variant: 'primary' as const } : {})}
          onPress={() => (one.key === 'auto' ? toggleAuto($) : compact($))}
        />
        {one.hint !== null && <Text dimColor>{` ${one.hint}`}</Text>}
      </Box>
    )
    const lines = layout.lines.map((line, index) => (
      <Box key={layout.lines.length > 1 ? `zeile-${index + 1}` : 'zeile'} overflow="hidden">
        {line.segments.flatMap(segment)}
        {line.controls.map(control)}
      </Box>
    ))

    let warning: RenderElement | null = null
    if (isWarning) {
      const form = warningLayout(remainingMs(s, cfg) ?? 0, e.props.bodyColumns)
      warning = (
        <Box key="warnung" overflow="hidden">
          <Text color="error" bold wrap="truncate">
            {form.text}
          </Text>
          <Box marginLeft={2} flexShrink={0}>
            <Button
              key="ping-once"
              label={form.pingLabel}
              hotkey="p"
              variant="primary"
              onPress={() => pingOnce($)}
            />
          </Box>
          <Box marginLeft={2} flexShrink={0}>
            <Button key="compact" label={form.compactLabel} hotkey="k" onPress={() => compact($)} />
          </Box>
        </Box>
      )
    }

    return (
      <Box flexDirection="column">
        {others}
        {warning}
        <Box key="statuszeile" flexDirection="column">
          {lines}
        </Box>
      </Box>
    )
  })
}

function toneProps(tone: Tone): { color?: string; dimColor?: boolean } {
  if (tone === 'plain') return {}
  if (tone === 'dim') return { dimColor: true }
  return { color: tone }
}

function asSaved(value: unknown): SavedSession | undefined {
  if (typeof value !== 'object' || value === null) return undefined
  const saved = value as Partial<SavedSession>
  return typeof saved.savedAt === 'number' && typeof saved.auto === 'object' && saved.auto !== null
    ? (saved as SavedSession)
    : undefined
}