/** @jsxImportSource @opentui/solid */
import type { TuiPlugin, TuiPluginModule } from "@opencode-ai/plugin/tui"
import { createMemo, createSignal } from "solid-js"
import { writeFileSync, appendFileSync } from "fs"

const DEBUG_LOG = "/tmp/oc-ls-stats-debug.log"
function debug(...args: unknown[]) {
  try {
    const line = `${args.map(a => typeof a === "object" ? JSON.stringify(a) : String(a)).join(" ")}\n`
    appendFileSync(DEBUG_LOG, line)
  } catch {}
}

type StreamSample = {
  at: number
  tokens: number
}

const STREAM_WINDOW_MS = 5_000
const LIVE_STALE_MS = 1_500
const SINGLE_SAMPLE_MS = 1_000
const SLOT_POLL_MS = 500
type MessageTiming = {
  sessionID: string
  requestStartAt: number
  firstResponseAt?: number
  firstTokenAt?: number
  lastTokenAt?: number
  lastToolCallAt?: number
}

type SessionAverage = {
  totalTokens: number
  totalDurationMs: number
  totalTtftMs: number
  messageCount: number
}

type TrackerState = {
  streamSamplesBySession: Record<string, StreamSample[]>
  messageTimingByID: Record<string, MessageTiming>
  sessionAverageByID: Record<string, SessionAverage>
  lastPrefillRate: number
  lastGeneratedTps: number
  isPrefilling: boolean
  prefillSlotId: number | null
  prefillCapturedTokens: number | null
  prefillStartAt: number
  isGenerating: boolean
  generateSlotId: number | null
  generateCapturedDecoded: number | null
  generateStartAt: number
  generateMaxDecoded: number
}

function estimateStreamTokens(delta: string) {
  return Math.max(1, Math.ceil(Buffer.byteLength(delta, "utf8") / 5))
}

function parseJSONC(text: string): unknown {
  let result = ""
  let inString = false
  let escape = false
  let i = 0

  while (i < text.length) {
    const ch = text[i]

    if (escape) {
      result += ch
      escape = false
      i++
      continue
    }

    if (ch === "\\") {
      result += ch
      escape = true
      i++
      continue
    }

    if (ch === '"') {
      result += ch
      inString = !inString
      i++
      continue
    }

    if (inString) {
      result += ch
      i++
      continue
    }

    if (ch === "/" && text[i + 1] === "/") {
      while (i < text.length && text[i] !== "\n") i++
      continue
    }

    if (ch === "/" && text[i + 1] === "*") {
      i += 2
      while (i < text.length - 1 && !(text[i] === "*" && text[i + 1] === "/")) i++
      i += 2
      result += " "
      continue
    }

    if (ch === ",") {
      const next = text.indexOf(/[\s}\]]/, i + 1)
      if (next !== -1 && /[\s}\]]/.test(text[next])) {
        i++
        continue
      }
    }

    result += ch
    i++
  }

  return JSON.parse(result)
}

function stripBaseUrlPath(baseURL: string): string {
  try {
    const url = new URL(baseURL)
    return `${url.origin}`
  } catch {
    return baseURL.replace(/\/v1(\/.*)?$/, "").replace(/\/$/, "") || baseURL
  }
}

function extractProviderUrls(config: unknown): string[] {
  if (!config || typeof config !== "object") return []
  const obj = config as Record<string, unknown>
  const provider = obj.provider
  if (!provider || typeof provider !== "object") return []

  const urls: string[] = []
  for (const [_key, val] of Object.entries(provider as Record<string, unknown>)) {
    if (val && typeof val === "object") {
      const opts = (val as Record<string, unknown>).options
      if (opts && typeof opts === "object") {
        const baseURL = (opts as Record<string, unknown>).baseURL || (opts as Record<string, unknown>).base_url
        if (typeof baseURL === "string" && baseURL.includes("localhost")) {
          urls.push(stripBaseUrlPath(baseURL))
        }
      }
    }
  }
  return urls
}

async function loadConfigUrls(api: Parameters<TuiPlugin>[0]): Promise<string[]> {
  try {
    const configPath = api.state.path.config
    const resp = await fetch(`file://${configPath}`)
    if (!resp.ok) return []
    const text = await resp.text()
    const config = parseJSONC(text)
    return extractProviderUrls(config)
  } catch {
    return []
  }
}

async function fetchSlots(baseUrl: string, model?: string): Promise<unknown | null> {
  try {
    const url = model ? `${baseUrl}/slots?model=${encodeURIComponent(model)}` : `${baseUrl}/slots`
    const resp = await fetch(url)
    if (!resp.ok) return null
    return await resp.json()
  } catch {
    return null
  }
}

function formatPps(value: number) {
  if (!Number.isFinite(value) || value <= 0) return undefined
  return `${Math.round(value)}`
}

function formatTps(value: number) {
  if (!Number.isFinite(value) || value <= 0) return undefined
  return `${Math.round(value)}`
}

function formatAvg(value: number) {
  if (!Number.isFinite(value) || value <= 0) return undefined
  if (value >= 100) return `${Math.round(value)}`
  return `${value.toFixed(1)}`
}

function pad4(s: string | undefined) {
  if (!s) return "----"
  return s.padStart(4, " ")
}

function activeDurationMs(samples: StreamSample[], tailAt?: number) {
  if (samples.length === 0) return 0
  if (samples.length === 1) {
    const tailDuration = tailAt ? Math.max(0, tailAt - samples[0].at) : SINGLE_SAMPLE_MS
    return Math.min(Math.max(tailDuration, 250), SINGLE_SAMPLE_MS)
  }

  let duration = 0
  for (let i = 1; i < samples.length; i++) {
    duration += Math.max(0, samples[i].at - samples[i - 1].at)
  }

  if (tailAt) {
    duration += Math.max(0, tailAt - samples[samples.length - 1].at)
  }

  return Math.max(duration, SINGLE_SAMPLE_MS)
}

function SessionPromptRight(props: {
  api: Parameters<TuiPlugin>[0]
  sessionID: string
  tracker: TrackerState
  version: () => number
  clock: () => number
}) {
  const liveTps = createMemo(() => {
    props.version()
    props.clock()
    const rate = props.tracker.lastGeneratedTps
    if (!Number.isFinite(rate) || rate <= 0) return undefined
    return rate
  })

  const sessionAverage = createMemo(() => {
    props.version()
    const totals = props.tracker.sessionAverageByID[props.sessionID]
    if (!totals || totals.totalTokens <= 0 || totals.totalDurationMs <= 0) return undefined
    return totals.totalTokens / (totals.totalDurationMs / 1000)
  })

  const prefillRate = createMemo(() => {
    props.version()
    if (!props.tracker.isPrefilling) return undefined
    const rate = props.tracker.lastPrefillRate
    if (!Number.isFinite(rate) || rate <= 0) return undefined
    return rate
  })

   const text = createMemo(() => {
    const pp = prefillRate()
    const tps = liveTps()
    const avg = sessionAverage()
    const parts = [`PP ${pad4(pp ? formatPps(pp) : undefined)}`, `TPS ${pad4(tps ? formatTps(tps) : undefined)}`, `AVG ${pad4(avg ? formatAvg(avg) : undefined)}`]
    return parts.join(" | ")
  })

  return <>{text() ? <text fg={props.api.theme.current.textMuted}>{text()}</text> : null}</>
}

const tui: TuiPlugin = async (api) => {
  console.log("[oc-ls-stats] TUI plugin loaded!")
  const tracker: TrackerState = {
    streamSamplesBySession: {},
    messageTimingByID: {},
    sessionAverageByID: {},
    lastPrefillRate: 0,
    lastGeneratedTps: 0,
    isPrefilling: false,
    prefillSlotId: null,
    prefillCapturedTokens: null,
    prefillStartAt: 0,
    isGenerating: false,
    generateSlotId: null,
    generateCapturedDecoded: null,
    generateStartAt: 0,
    generateMaxDecoded: 0,
  }
  const [version, setVersion] = createSignal(0)
  const [clock, setClock] = createSignal(Date.now())

  const bump = () => setVersion((value) => value + 1)

  const pruneSamples = (now = Date.now()) => {
    let changed = false

    for (const [sessionID, samples] of Object.entries(tracker.streamSamplesBySession)) {
      const next = samples.filter((sample) => now - sample.at <= STREAM_WINDOW_MS)
      if (next.length !== samples.length) {
        changed = true
        if (next.length > 0) tracker.streamSamplesBySession[sessionID] = next
        else delete tracker.streamSamplesBySession[sessionID]
      }
    }

    if (changed) bump()
  }

  const clearLiveSamples = (sessionID: string) => {
    if (!tracker.streamSamplesBySession[sessionID]?.length) return
    delete tracker.streamSamplesBySession[sessionID]
    bump()
  }

  const appendSample = (sessionID: string, messageID: string, sample: StreamSample) => {
    const now = sample.at
    tracker.streamSamplesBySession[sessionID] = [
      ...(tracker.streamSamplesBySession[sessionID] ?? []).filter((item) => now - item.at <= STREAM_WINDOW_MS),
      sample,
    ]
    const timing = tracker.messageTimingByID[messageID]
    if (timing) {
      tracker.messageTimingByID[messageID] = timing.firstTokenAt
        ? { ...timing, lastTokenAt: now }
        : {
            ...timing,
            firstResponseAt: timing.firstResponseAt ?? now,
            firstTokenAt: now,
            lastTokenAt: now,
          }
    }
    bump()
  }

  let llamaServerUrls: string[] = ["http://localhost:8080"]
  let llamaServerModel: string | undefined

  const pollMetrics = async () => {
    if (!llamaServerUrls.length) return

    let model = llamaServerModel
    if (!model) {
      const routeSessionID = api.route.current.params?.sessionID
      if (routeSessionID) {
        const session = api.state.session.get(routeSessionID)
        if (session?.model?.id) {
          model = session.model.id
        }
      }
    }

    let anyPrefilling = false
    let anyGenerating = false
    let genEnded = false

    for (const baseUrl of llamaServerUrls) {
      const slots = await fetchSlots(baseUrl, model)
      const slotList = Array.isArray(slots) ? slots : slots ? Object.values(slots) : []

      let prefillSlot: any = null
      let generateSlot: any = null
      for (const slot of slotList) {
        if (!slot?.is_processing) continue
        if (slot.next_token?.[0]?.n_remain === -1) {
          prefillSlot = slot
          break
        }
        if (slot.next_token?.[0]?.n_decoded > 0) {
          generateSlot = slot
        }
      }

      if (prefillSlot) {
        anyPrefilling = true
        const slotId = prefillSlot.id
        const now = Date.now()
        const nt = prefillSlot.n_prompt_tokens ?? 0

        if (tracker.prefillSlotId === slotId) {
          if (tracker.prefillCapturedTokens !== null) {
            const dt = (now - tracker.prefillStartAt) / 1000
            const delta = nt - tracker.prefillCapturedTokens
            if (dt > 0 && delta > 0) {
              tracker.lastPrefillRate = delta / dt
            }
          }
        } else {
          tracker.prefillSlotId = slotId
          tracker.prefillCapturedTokens = nt
          tracker.prefillStartAt = now
          tracker.lastPrefillRate = 0
        }
      } else {
        if (tracker.isPrefilling) {
          tracker.prefillSlotId = null
          tracker.prefillCapturedTokens = null
          tracker.prefillStartAt = 0
        }
      }

      if (generateSlot) {
        anyGenerating = true
        const slotId = generateSlot.id
        const now = Date.now()
        const nd = generateSlot.next_token?.[0]?.n_decoded ?? 0

        if (nd > tracker.generateMaxDecoded) {
          const isNew = tracker.generateMaxDecoded === 0
          const delta = nd - tracker.generateMaxDecoded
          const dt = (now - tracker.generateStartAt) / 1000
          if (delta > 0 && dt > 0 && !isNew) {
            tracker.lastGeneratedTps = delta / dt
            debug(`gen slot=${slotId} nd=${nd} max=${tracker.generateMaxDecoded} delta=${delta} dt=${dt.toFixed(2)} tps=${tracker.lastGeneratedTps.toFixed(1)}`)
          }
          tracker.generateMaxDecoded = nd
          tracker.generateCapturedDecoded = nd
          tracker.generateStartAt = now
          tracker.generateSlotId = slotId
        } else if (nd < tracker.generateMaxDecoded) {
          tracker.generateCapturedDecoded = nd
          tracker.generateStartAt = now
          tracker.lastGeneratedTps = 0
          debug(`gen slot=${slotId} nd=${nd} max=${tracker.generateMaxDecoded} counter reset`)
        } else {
          debug(`gen slot=${slotId} nd=${nd} max=${tracker.generateMaxDecoded} stalled`)
          tracker.lastGeneratedTps = 0
          genEnded = true
        }
      } else {
        if (tracker.isGenerating) {
          tracker.generateSlotId = null
          tracker.generateCapturedDecoded = null
          tracker.generateStartAt = 0
          tracker.generateMaxDecoded = 0
          tracker.lastGeneratedTps = 0
          genEnded = true
          debug(`gen no slot found ending`)
        }
      }

      tracker.isPrefilling = anyPrefilling
      tracker.isGenerating = anyGenerating
    }

    if (tracker.lastPrefillRate > 0 || tracker.lastGeneratedTps > 0 || genEnded) {
      bump()
    }
  }

  const onDelta = api.event.on("message.part.delta", (evt) => {
    if (evt.properties.field !== "text") return
    const parts = api.state.part(evt.properties.messageID)
    const part = parts.find((item) => item.id === evt.properties.partID)
    if (!part) return
    if (part.type !== "text" && part.type !== "reasoning") return
    appendSample(evt.properties.sessionID, evt.properties.messageID, {
      at: Date.now(),
      tokens: estimateStreamTokens(evt.properties.delta),
    })
  })

  const onMessage = api.event.on("message.updated", (evt) => {
    if (evt.properties.info.role !== "assistant") return

    if (!evt.properties.info.time.completed) {
      const existing = tracker.messageTimingByID[evt.properties.info.id]
      tracker.messageTimingByID[evt.properties.info.id] = {
        sessionID: evt.properties.sessionID,
        requestStartAt: evt.properties.info.time.created,
        firstResponseAt: existing?.firstResponseAt,
        firstTokenAt: existing?.firstTokenAt,
        lastTokenAt: existing?.lastTokenAt,
        lastToolCallAt: existing?.lastToolCallAt,
      }
      bump()
      return
    }

    const timing = tracker.messageTimingByID[evt.properties.info.id]
    if (timing?.sessionID === evt.properties.sessionID && typeof timing.firstResponseAt === "number") {
      const totalTokens = evt.properties.info.tokens.output + evt.properties.info.tokens.reasoning
      const endAt =
        evt.properties.info.finish === "tool-calls"
          ? timing.lastToolCallAt
          : evt.properties.info.time.completed
      const durationMs = typeof endAt === "number" ? Math.max(endAt - timing.firstResponseAt, 1) : undefined
      const ttftMs = Math.max(timing.firstResponseAt - timing.requestStartAt, 0)
      if (totalTokens > 0 && durationMs) {
        const totals = tracker.sessionAverageByID[evt.properties.sessionID] ?? {
          totalTokens: 0,
          totalDurationMs: 0,
          totalTtftMs: 0,
          messageCount: 0,
        }
        tracker.sessionAverageByID[evt.properties.sessionID] = {
          totalTokens: totals.totalTokens + totalTokens,
          totalDurationMs: totals.totalDurationMs + durationMs,
          totalTtftMs: totals.totalTtftMs + ttftMs,
          messageCount: totals.messageCount + 1,
        }
      }
    }
    delete tracker.messageTimingByID[evt.properties.info.id]
    pruneSamples(evt.properties.info.time.completed)
    bump()
  })

  const onPart = api.event.on("message.part.updated", (evt) => {
    if (evt.properties.part.type !== "tool") return
    if (
      evt.properties.part.state.status === "running" ||
      evt.properties.part.state.status === "completed" ||
      evt.properties.part.state.status === "error"
    ) {
      clearLiveSamples(evt.properties.sessionID)
    }
    const timing = tracker.messageTimingByID[evt.properties.part.messageID]
    if (!timing) return
    if (evt.properties.part.state.status === "pending") {
      tracker.messageTimingByID[evt.properties.part.messageID] = {
        ...timing,
        firstResponseAt: timing.firstResponseAt ?? evt.properties.time,
      }
      bump()
      return
    }
    if (evt.properties.part.state.status !== "running") return
    tracker.messageTimingByID[evt.properties.part.messageID] = {
      ...timing,
      lastToolCallAt: evt.properties.part.state.time.start,
    }
    bump()
  })

  const timer = setInterval(() => {
    setClock(Date.now())
    pruneSamples()
  }, 1000)

  const slotTimer = setInterval(async () => {
    await pollMetrics()
  }, SLOT_POLL_MS)

  api.lifecycle.onDispose(() => {
    onDelta()
    onMessage()
    onPart()
    clearInterval(timer)
    clearInterval(slotTimer)
  })

  try {
    if (!llamaServerUrls.length) {
      llamaServerUrls = await loadConfigUrls(api)
    }
  } catch (e) {
    console.error("[oc-ls-stats] loadConfigUrls error:", e)
  }

  api.slots.register({
    slots: {
      session_prompt_right(_ctx, value) {
        return <SessionPromptRight api={api} sessionID={value.session_id} tracker={tracker} version={version} clock={clock} />
      },
    },
  })
}

const plugin: TuiPluginModule & { id: string } = {
  id: "@troed/oc-ls-stats",
  tui,
}

export default plugin