import { atom, read, update } from 'claude-code'
import type { EngineInterface, Register, RenderChildren } from 'claude-code'

import type { Correction, Corrections, Exception, RuleRef } from '../types'

const PANE = 'straion-trail'
const TITLE = 'Straion'
const POLL_MS = 1500
const SAFE_SESSION_ID = /^[A-Za-z0-9_-]{1,128}$/
const SIGNET_PX = 20
const TERMINAL_SIGNET_COLUMNS = 4
const TERMINAL_SIGNET_ROWS = 2
const SIGNET_SVG =
  '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 197 197"><style>path{fill:#191928}@media (prefers-color-scheme: dark){path{fill:#D4CFC3}}</style><path d="M0 197H196.416V0H0V197ZM105.579 167.462V78.8058C129.993 78.8058 149.773 98.6506 149.773 123.131C149.773 147.611 129.987 167.456 105.579 167.456V167.462ZM147.324 29.5268V73.875C122.899 73.875 103.113 54.0186 103.113 29.5268H147.324ZM49.0983 167.473V123.131C73.5232 123.131 93.3092 142.987 93.3092 167.473H49.0983ZM90.8425 29.55V118.327C66.4002 118.327 46.5795 98.4536 46.5795 73.9329C46.5795 49.4122 66.3944 29.5442 90.8425 29.5442V29.55Z"/></svg>'

const trail = atom({ plugin: 'straion-trail', key: 'trail' } as const, null)
const followed = atom({ plugin: 'straion-trail', key: 'followed' } as const, null)
const showCompliant = atom({ plugin: 'straion-trail', key: 'showCompliant' } as const, false)
const showExceptions = atom({ plugin: 'straion-trail', key: 'showExceptions' } as const, false)
const isClosed = atom({ plugin: 'straion-trail', key: 'isClosed' } as const, false)

/** The parts of `straion audit view --json` (the CLI's `AuditModel`) the pane reads. */
type AuditRule = {
  shortcode: string
  statement: string
  state: string
  decision: string | null
  reason: string | null
  humanConfirmed: boolean
  priorVerdicts: { state: string; reason: string | null }[]
}

type AuditModel = { rules: AuditRule[]; awaitingApproval?: string[] }

function firstSentence(text: string): string {
  return text.split(/(?<=\.)\s/)[0]?.replace(/\.$/, '') ?? text
}

function findFixReason(rule: AuditRule): string | null {
  if (rule.state === 'fix-required') return rule.reason
  const fixVerdict = [...rule.priorVerdicts].reverse().find(prior => prior.state === 'fix-required')
  return fixVerdict?.reason ?? null
}

function wasCorrected(rule: AuditRule): boolean {
  return rule.decision === 'fix' || rule.priorVerdicts.some(prior => prior.state === 'fix-required')
}

function isApprovedException(rule: AuditRule): boolean {
  return (rule.state === 'accepted' || rule.state === 'not-applicable') && rule.humanConfirmed
}

function deriveCorrections(sessionId: string, model: AuditModel): Corrections {
  const ref = (rule: AuditRule): RuleRef => ({ shortcode: rule.shortcode, rule: firstSentence(rule.statement) })
  const awaiting = new Set(model.awaitingApproval ?? [])
  const corrections: Correction[] = model.rules
    .filter(wasCorrected)
    .map(rule => ({ ...ref(rule), fix: findFixReason(rule), isFixed: rule.state === 'compliant' }))
    .sort((a, b) => Number(a.isFixed) - Number(b.isFixed))

  return {
    sessionId,
    isFound: true,
    ruleCount: model.rules.length,
    awaiting: model.rules.filter(rule => awaiting.has(rule.shortcode)).map(ref),
    corrections,
    exceptions: model.rules
      .filter(isApprovedException)
      .map((rule): Exception => ({ ...ref(rule), kind: rule.state === 'accepted' ? 'accepted' : 'not-applicable' })),
    compliant: model.rules.filter(rule => rule.state === 'compliant' && !wasCorrected(rule)).map(ref),
  }
}

function summarize(trail: Corrections): string {
  const count = trail.corrections.length
  if (count === 0) return 'Straion kept everything on track'
  return count === 1 ? 'Straion corrected the agent once' : `Straion corrected the agent ${count} times`
}

let lastFileKey = ''
let lastPaneRows = 0
let canDrawImages = false
let isPrimed = false
const seen = new Set<string>()

function announce($: EngineInterface, next: Corrections): void {
  const moments = [
    ...next.awaiting.map(item => ({ key: `ask:${item.shortcode}`, text: `Straion needs your approval: ${item.rule}` })),
    ...next.corrections.map(item => ({
      key: `${item.isFixed ? 'fixed' : 'fix'}:${item.shortcode}`,
      text: `${item.isFixed ? 'Fixed' : 'Straion corrected'}: ${item.rule}`,
    })),
  ]
  for (const moment of moments) {
    if (isPrimed && !seen.has(moment.key)) $.ui.toast(moment.text)
    seen.add(moment.key)
  }
  isPrimed = true
}

/** Null when the CLI could not answer; the pane keeps what it last showed. */
async function readAuditModel($: EngineInterface, sessionId: string): Promise<AuditModel | null> {
  const ran = await $.process.run(['straion', 'audit', 'view', sessionId, '--json']).catch(() => null)
  if (ran === null || ran.exitCode !== 0) return null
  try {
    return JSON.parse(ran.stdout) as AuditModel
  } catch {
    return null
  }
}

async function refresh($: EngineInterface): Promise<void> {
  const sessionId = (await read($, followed)) ?? (await $.session.id())
  const home = await $.env.get('HOME')
  if (!home || !SAFE_SESSION_ID.test(sessionId)) return
  const path = `${home}/.straion/audit/${sessionId}.jsonl`
  const stat = await $.fs.stat(path).catch(() => undefined)
  const fileKey = stat ? `${path}:${stat.mtimeMs}:${stat.size}` : `${path}:missing`
  if (fileKey === lastFileKey) return

  if (!stat) {
    lastFileKey = fileKey
    await update($, trail, () => ({ sessionId, isFound: false, ruleCount: 0, awaiting: [], corrections: [], exceptions: [], compliant: [] }))
    $.ui.status(undefined)
    return
  }
  const model = await readAuditModel($, sessionId)
  if (model === null) return
  lastFileKey = fileKey
  const next = deriveCorrections(sessionId, model)
  announce($, next)
  await update($, trail, () => next)
  $.ui.status(next.corrections.length > 0 ? `Straion · ${next.corrections.length} corrected` : undefined)
  await fitPane($)
}

async function follow($: EngineInterface, sessionId: string | null): Promise<void> {
  await update($, followed, () => sessionId)
  lastFileKey = ''
  isPrimed = false
  seen.clear()
  await refresh($)
}

async function detectImageSupport($: EngineInterface): Promise<boolean> {
  const term = await $.env.get('TERM')
  const termProgram = await $.env.get('TERM_PROGRAM')
  const kittyWindow = await $.env.get('KITTY_WINDOW_ID')
  return term === 'xterm-kitty' || kittyWindow !== undefined || termProgram?.toLowerCase() === 'ghostty'
}

function countCorrectionLines(trail: Corrections): number {
  const awaiting = trail.awaiting.length * 2
  const corrections = trail.corrections.reduce(
    (lines, correction) => lines + (correction.isFixed || correction.fix === null ? 1 : 2),
    0,
  )
  return awaiting + corrections
}

/** The rows the compact terminal layout draws, so an inline pane opens exactly that tall. */
async function countPaneRows($: EngineInterface): Promise<number> {
  const current = await read($, trail)
  const headerRows = 2
  if (current === null || current.ruleCount === 0) return headerRows
  const toggles = [
    { count: current.exceptions.length, isOpen: await read($, showExceptions) },
    { count: current.compliant.length, isOpen: await read($, showCompliant) },
  ].filter(toggle => toggle.count > 0)
  const toggleRows = toggles.reduce((rows, toggle) => rows + 1 + (toggle.isOpen ? toggle.count : 0), 0)
  const bodyRows = countCorrectionLines(current) + toggleRows
  return bodyRows === 0 ? headerRows : headerRows + 1 + bodyRows
}

async function openPane($: EngineInterface): Promise<void> {
  lastPaneRows = await countPaneRows($)
  await $.ui.open({ id: PANE, title: TITLE, rows: lastPaneRows })
}

async function fitPane($: EngineInterface): Promise<void> {
  if (await read($, isClosed)) return
  if ((await countPaneRows($)) === lastPaneRows) return
  await openPane($)
}

export const register: Register = on => {
  on('session.start', async ($, e, next) => {
    await $.command.register({
      name: 'straion-trail',
      description: 'Show what Straion corrected. "close" hides it; a session id replays that session, "this" comes back.',
    })
    canDrawImages = await detectImageSupport($)
    await refresh($)
    $.clock.every(POLL_MS, () => refresh($))
    if (!(await read($, isClosed))) void openPane($)

    return next(e)
  })

  on('command.run', { command: 'straion-trail' }, async ($, e) => {
    const arg = e.args.trim()
    if (arg === 'close') {
      await update($, isClosed, () => true)
      await $.ui.close({ id: PANE })
      return { text: 'Straion pane closed. /straion-trail opens it again.' }
    }
    if (arg !== '' && arg !== 'this' && !SAFE_SESSION_ID.test(arg)) {
      return { text: `"${arg}" is not a session id.` }
    }
    if (arg !== '') await follow($, arg === 'this' ? null : arg)
    await update($, isClosed, () => false)
    await openPane($)

    return { text: 'Straion pane opened.' }
  })

  on('ui.close', async ($, e, next) => {
    if (e.id === PANE && e.origin.kind === 'person') await update($, isClosed, () => true)

    return next(e)
  })

  on('tool.call', async ($, e, next) => {
    const ran = await next(e)
    await refresh($)

    return ran
  })

  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
    const { Box, Text, Button } = $.ui.resolve(e)
    const current = await read($, trail)
    const isShowingCompliant = await read($, showCompliant)
    const isShowingExceptions = await read($, showExceptions)
    const width = Math.max(20, e.props.bodyColumns ?? 60)
    const isCompact = e.surface === 'terminal'
    const itemGap = isCompact ? 0 : 1
    const sectionGap = isCompact ? 1 : 2

    const mark = (icon: string, color: string) => {
      if (e.surface !== 'desktop') return <Text color={color}>{icon}</Text>
      const { Svg } = $.ui.resolve(e)
      return <Svg source={SIGNET_SVG} alt="Straion" width={SIGNET_PX} height={SIGNET_PX} />
    }

    const terminalMark = () => {
      if (e.surface !== 'terminal' || !canDrawImages) return null
      const { Image } = $.ui.resolve(e)
      return (
        <Box width={TERMINAL_SIGNET_COLUMNS + 1} flexShrink={0}>
          <Image
            source={{ file: `${$.plugin.root}/assets/signet.png`, format: 'png' }}
            columns={TERMINAL_SIGNET_COLUMNS}
            rows={TERMINAL_SIGNET_ROWS}
            alt=" "
          />
        </Box>
      )
    }

    const header = (icon: string, color: string, title: string, subline: string) =>
      isCompact ? (
        <Box flexDirection="row">
          {terminalMark()}
          <Box flexDirection="column" flexGrow={1}>
            <Text bold wrap="truncate-end">{title}</Text>
            <Text dimColor wrap="truncate-end">
              <Text color={color}>{icon}</Text> {subline}
            </Text>
          </Box>
        </Box>
      ) : (
        <Box flexDirection="column" alignItems="center" paddingTop={2}>
          {mark(icon, color)}
          <Box marginTop={1}>
            <Text bold>{title}</Text>
          </Box>
          <Text dimColor>
            <Text color={color}>{icon} </Text>
            {subline}
          </Text>
        </Box>
      )

    const item = (marker: string, color: string, title: string, detail: string | null, isQuiet: boolean) => (
      <Box flexDirection="row" marginTop={itemGap}>
        <Box width={2} flexShrink={0}>
          <Text color={color}>{marker}</Text>
        </Box>
        <Box flexDirection="column" flexGrow={1}>
          <Text dimColor={isQuiet} wrap="truncate-end">{title}</Text>
          {detail !== null ? <Text dimColor wrap="truncate-end">{detail}</Text> : null}
        </Box>
      </Box>
    )

    const collapsible = (key: string, isOpen: boolean, label: string, onPress: () => unknown, rows: RenderChildren[]) => (
      <Box key={key} flexDirection="column">
        <Button key={`toggle-${key}`} plain label={isOpen ? `Hide ${label}` : `Show ${rows.length} ${label}`} onPress={onPress} />
        {isOpen ? <Box flexDirection="column">{rows}</Box> : null}
      </Box>
    )

    if (current === null || current.ruleCount === 0) {
      return (
        <Box flexDirection="column" width={width}>
          {header('◇', 'green', 'Straion is watching', 'Corrections show up here')}
        </Box>
      )
    }

    const fixed = current.corrections.filter(correction => correction.isFixed).length
    const inProgress = current.corrections.length - fixed
    const hero =
      current.awaiting.length > 0
        ? { icon: '?', color: 'yellow' }
        : inProgress > 0
          ? { icon: '●', color: 'red' }
          : { icon: '✓', color: 'green' }
    const subline =
      current.corrections.length === 0
        ? `${current.ruleCount} rules checked`
        : inProgress > 0
          ? `${fixed} fixed · ${inProgress} in progress`
          : `all fixed · ${current.ruleCount} rules checked`

    const listed = [
      ...current.awaiting.map(rule => (
        <Box key={`ask-${rule.shortcode}`} flexDirection="column">
          {item('?', 'yellow', 'Your approval is needed', rule.rule, false)}
        </Box>
      )),
      ...current.corrections.map(correction => (
        <Box key={correction.shortcode} flexDirection="column">
          {correction.isFixed
            ? item('✓', 'green', correction.rule, null, false)
            : item('●', 'red', correction.rule, correction.fix, false)}
        </Box>
      )),
    ]
    const sections = [
      listed.length > 0 ? <Box key="listed" flexDirection="column">{listed}</Box> : null,
      current.exceptions.length > 0
        ? collapsible(
            'exceptions',
            isShowingExceptions,
            'approved exceptions',
            () => update($, showExceptions, isShown => !isShown).then(() => fitPane($)),
            current.exceptions.map(rule => (
              <Box key={`ex-${rule.shortcode}`} flexDirection="column">
                {item(
                  '–',
                  'gray',
                  rule.rule,
                  isCompact ? null : rule.kind === 'accepted' ? 'accepted as is' : 'not applicable here',
                  true,
                )}
              </Box>
            )),
          )
        : null,
      current.compliant.length > 0
        ? collapsible(
            'compliant',
            isShowingCompliant,
            'compliant rules',
            () => update($, showCompliant, isShown => !isShown).then(() => fitPane($)),
            current.compliant.map(rule => (
              <Box key={`ok-${rule.shortcode}`} flexDirection="column">
                {item('✓', 'gray', rule.rule, null, true)}
              </Box>
            )),
          )
        : null,
    ].filter(section => section !== null)

    return (
      <Box flexDirection="column" width={width}>
        {header(hero.icon, hero.color, summarize(current), subline)}
        <Box flexDirection="column" paddingX={isCompact ? 0 : 1}>
          {sections.map((section, index) => (
            <Box key={`section-${index}`} flexDirection="column" marginTop={index === 0 ? sectionGap : isCompact ? 0 : sectionGap}>
              {section}
            </Box>
          ))}
        </Box>
      </Box>
    )
  })
}