/// <reference path="./ambient-modules.d.ts" />
import type { ExtensionAPI, ExtensionCommandContext } from '@earendil-works/pi-coding-agent'
import { Type } from 'typebox'
import { selectOperationsFromUI } from './reverse-last-picker'
import type { SessionEntryLike, UndoOperation } from './reverse-last-state'

type UndoContext = {
  cwd: string
  sessionManager: {
    getSessionId(): string
    getLeafId(): string | null
    getBranch?: () => SessionEntryLike[]
    getEntries(): SessionEntryLike[]
  }
}

type UndoResult = { steps: number; restored: string[] }

export function registerReverseLastCommands(
  pi: ExtensionAPI,
  dependencies: {
    pending: Map<string, UndoOperation>
    getSelectionCandidates: (ctx: UndoContext) => UndoOperation[]
    undoChanges: (ctx: UndoContext, count: number) => UndoResult
    undoSelectedChanges: (ctx: UndoContext, selectedIds: Set<string>) => UndoResult
  }
): void {
  const { pending, getSelectionCandidates, undoChanges, undoSelectedChanges } = dependencies

  pi.registerCommand('undo', {
    description: 'Undo last Pi write/edit file changes. Usage: /undo [count]',
    handler: async (args: string, ctx: ExtensionCommandContext) => {
      await ctx.waitForIdle()
      if (pending.size > 0) {
        ctx.ui.notify(
          '/undo is temporarily unavailable while file changes are still being captured',
          'warning'
        )
        return
      }

      const trimmed = args.trim()
      if (!trimmed && ctx.hasUI) {
        const candidates = getSelectionCandidates(ctx)
        if (candidates.length === 0) {
          ctx.ui.notify('Nothing to undo in this session yet.', 'info')
          return
        }

        const selectedIds = await selectOperationsFromUI(ctx, candidates)
        if (selectedIds === undefined) {
          ctx.ui.notify('Undo cancelled.', 'info')
          return
        }
        if (selectedIds.size === 0) {
          ctx.ui.notify('No changes selected.', 'warning')
          return
        }

        const result = undoSelectedChanges(ctx, selectedIds)
        if (result.steps === 0) {
          ctx.ui.notify('Nothing to undo in this session yet.', 'info')
          return
        }
        ctx.ui.notify(formatNotification('selected change', result), 'info')
        return
      }

      const count = trimmed ? Number.parseInt(trimmed, 10) : 1
      if (!Number.isFinite(count) || count <= 0) {
        ctx.ui.notify('Usage: /undo [count]  (count must be a positive integer)', 'warning')
        return
      }

      const result = undoChanges(ctx, count)
      if (result.steps === 0) {
        ctx.ui.notify('Nothing to undo in this session yet.', 'info')
        return
      }
      ctx.ui.notify(formatNotification('change', result), 'info')
    }
  })

  pi.registerTool({
    name: 'undo_files',
    label: 'Undo Files',
    description: 'Undo the most recent write/edit file changes captured in this session',
    parameters: Type.Object({
      count: Type.Optional(
        Type.Number({
          minimum: 1,
          maximum: 20,
          description: 'How many captured change steps to undo'
        })
      )
    }),
    async execute(
      _toolCallId: string,
      params: { count?: number },
      _signal: unknown,
      _onUpdate: unknown,
      ctx
    ) {
      if (pending.size > 0) {
        throw new Error(
          'undo_files is temporarily unavailable while file changes are still being captured'
        )
      }

      const count =
        typeof params.count === 'number' ? Math.max(1, Math.min(20, Math.trunc(params.count))) : 1
      const result = undoChanges(ctx, count)
      if (result.steps === 0) {
        return {
          content: [{ type: 'text', text: 'Nothing to undo in this session yet.' }],
          details: result
        }
      }

      return {
        content: [{ type: 'text', text: formatNotification('change', result, 20) }],
        details: result
      }
    }
  })
}

function formatNotification(kind: string, result: UndoResult, maxFiles = 12): string {
  const suffix =
    result.restored.length > maxFiles ? `\n...and ${result.restored.length - maxFiles} more` : ''
  return `Reverted ${result.steps} ${kind}${result.steps === 1 ? '' : 's'}.\n${result.restored.slice(0, maxFiles).join('\n')}${suffix}`
}