/// <reference path="./ambient-modules.d.ts" />
import * as path from 'node:path'
import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
import { normalizeTimestampMs } from '@firstpick/pi-utils'
import { isToolCallEventType } from '@earendil-works/pi-coding-agent'
import {
  adoptUntrackedDiskEdits,
  applyTreeFileCheckout,
  findOutsideDeletions,
  findUntrackedDiskEdits,
  plannedTreeFileChanges,
  loadState,
  pruneStalePins,
  readTextFile,
  rebuildStackFromHistory,
  resolveToolPath,
  saveState,
  type SessionEntryLike,
  type SessionUndoState,
  type UndoOperation
} from './reverse-last-state.ts'
import {
  captureTurnLeafId,
  currentSessionLeafId,
  ensureLegacyHistory,
  recordCapturedOperation,
  snapshotChangedFiles,
  undoRecentOperations,
  undoSelectedOperations
} from './reverse-last-capture.ts'
import { registerReverseLastCommands } from './reverse-last-commands.ts'

const EXTENSION_ID = 'reverse-last'

function toolCallIdFromEvent(event: { toolCallId?: unknown }): string | undefined {
  return typeof event.toolCallId === 'string' ? event.toolCallId : undefined
}

function reverseOperations(operations: UndoOperation[]): UndoOperation[] {
  const reversed: UndoOperation[] = []
  for (let index = operations.length - 1; index >= 0; index--) {
    const operation = operations[index]
    if (operation) {
      reversed.push(operation)
    }
  }
  return reversed
}

export default function reverseLastExtension(pi: ExtensionAPI) {
  const pending = new Map<string, UndoOperation>()

  let activeSessionId = ''
  let turnLeafId: string | undefined
  let state: SessionUndoState = { stack: [], history: [] }

  const reloadForSession = (
    sessionId: string,
    sessionManager?: { getBranch?: () => SessionEntryLike[]; getEntries(): SessionEntryLike[] }
  ) => {
    activeSessionId = sessionId
    pending.clear()
    turnLeafId = undefined
    state = loadState(sessionId)
    if (sessionManager) {
      rebuildStackFromHistory(state, sessionManager)
    }
    pruneStalePins(state)
  }

  pi.on('session_start', async (_event, ctx) => {
    reloadForSession(ctx.sessionManager.getSessionId(), ctx.sessionManager)
  })

  pi.on('turn_end', async () => {
    turnLeafId = undefined
  })

  pi.on('session_tree', async (event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId()
    const sameSession = Boolean(activeSessionId) && sessionId === activeSessionId
    const previousStack = sameSession ? [...state.stack] : []

    if (sameSession) {
      rebuildStackFromHistory(state, ctx.sessionManager)
      pruneStalePins(state)
    } else {
      reloadForSession(sessionId, ctx.sessionManager)
    }

    if (sameSession && pending.size === 0) {
      const deletionPaths = findOutsideDeletions(previousStack, state.stack, state.pins)
      let allowOutsideDeletions = false
      if (deletionPaths.length > 0) {
        const preview = deletionPaths
          .slice(0, 8)
          .map((absPath) => path.relative(ctx.cwd, absPath) || absPath)
        const extra = deletionPaths.length > 8 ? `\n...and ${deletionPaths.length - 8} more` : ''
        const recreate = await ctx.ui.confirm(
          'Files deleted outside pi',
          `These captured files were deleted outside pi (shell/editor). Recreate them for this jump?\n${preview.join('\n')}${extra}\n\nAnswer No to leave them deleted; pi will not ask again until they change.`
        )
        if (recreate) {
          allowOutsideDeletions = true
        } else {
          const pins = state.pins ?? []
          state.pins = [...pins, ...deletionPaths.filter((absPath) => !pins.includes(absPath))]
          ctx.ui.notify(
            `Left ${deletionPaths.length} file${deletionPaths.length === 1 ? '' : 's'} deleted. pi will not ask about ${deletionPaths.length === 1 ? 'it' : 'them'} again until ${deletionPaths.length === 1 ? 'it' : 'they'} change.\n${preview.join('\n')}${extra}`,
            'info'
          )
        }
      }
      const dirty = findUntrackedDiskEdits(previousStack, state.stack, state.pins)
      let skipUntracked = true
      if (dirty.length > 0) {
        const preview = dirty
          .slice(0, 8)
          .map((absPath) => path.relative(ctx.cwd, absPath) || absPath)
        const extra = dirty.length > 8 ? `\n...and ${dirty.length - 8} more` : ''
        const overwrite = await ctx.ui.confirm(
          'Untracked file edits',
          `These files changed outside captured write/edit snapshots. Overwrite them for /tree?\n${preview.join('\n')}${extra}`
        )
        skipUntracked = !overwrite
        if (skipUntracked) {
          adoptUntrackedDiskEdits(state.stack, dirty)
          adoptUntrackedDiskEdits(
            state.history.flatMap((record) =>
              record.type === 'capture' ? [record.operation] : []
            ),
            dirty
          )
          ctx.ui.notify(
            `Kept current files and updated snapshots for ${dirty.length} file${dirty.length === 1 ? '' : 's'}. /tree will not ask about them again unless they change.\n${preview.join('\n')}${extra}`,
            'info'
          )
        }
      }
      const planned = plannedTreeFileChanges(previousStack, state.stack, {
        skipUntracked,
        pins: state.pins,
        allowOutsideDeletions
      })
      if (planned.length > 0) {
        const preview = planned.slice(0, 8).map((action) => {
          const label = path.relative(ctx.cwd, action.absPath) || action.absPath
          if (action.kind === 'delete') {
            return `${label} (delete)`
          }
          return readTextFile(action.absPath) === undefined ? `${label} (recreate)` : label
        })
        const extra = planned.length > 8 ? `\n...and ${planned.length - 8} more` : ''
        const restoreFiles = await ctx.ui.confirm(
          'Restore files?',
          `Restore captured files to this /tree node, or keep the current files and only move the conversation?\n${preview.join('\n')}${extra}`
        )
        if (!restoreFiles) {
          ctx.ui.notify('Jumped in the conversation only. Files were left unchanged.', 'info')
          if (activeSessionId) {
            saveState(activeSessionId, state)
          }
          return
        }
      }
      const restored = applyTreeFileCheckout(previousStack, state.stack, ctx.cwd, {
        skipUntracked,
        pins: state.pins,
        allowOutsideDeletions
      })
      if (restored.length > 0) {
        const suffix = restored.length > 12 ? `\n...and ${restored.length - 12} more` : ''
        ctx.ui.notify(
          `Restored ${restored.length} file${restored.length === 1 ? '' : 's'} for /tree.\n${restored.slice(0, 12).join('\n')}${suffix}`,
          'info'
        )
        pi.appendEntry(EXTENSION_ID, {
          type: 'tree-restore',
          at: Date.now(),
          oldLeafId: event.oldLeafId,
          newLeafId: event.newLeafId,
          count: restored.length,
          files: restored
        })
      }
    } else if (sameSession && pending.size > 0) {
      ctx.ui.notify(
        'Skipped file restore on /tree because a file change is still being captured',
        'warning'
      )
    }

    if (activeSessionId) {
      saveState(activeSessionId, state)
    }
  })

  pi.on('session_shutdown', async () => {
    if (!activeSessionId) {
      return
    }
    saveState(activeSessionId, state)
  })

  pi.on('tool_call', async (event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId()
    if (!activeSessionId) {
      reloadForSession(sessionId, ctx.sessionManager)
    }
    if (sessionId !== activeSessionId) {
      reloadForSession(sessionId, ctx.sessionManager)
    }

    if (isToolCallEventType('write', event)) {
      const absPath = resolveToolPath(ctx.cwd, String(event.input.path))
      const before = readTextFile(absPath)
      pending.set(event.toolCallId, {
        id: event.toolCallId,
        toolName: 'write',
        timestamp: Date.now(),
        files: [
          {
            absPath,
            previousExists: before !== undefined,
            previousContent: before
          }
        ]
      })
      return
    }

    if (isToolCallEventType('edit', event)) {
      const absPath = resolveToolPath(ctx.cwd, String(event.input.path))
      const before = readTextFile(absPath)
      pending.set(event.toolCallId, {
        id: event.toolCallId,
        toolName: 'edit',
        timestamp: Date.now(),
        files: [
          {
            absPath,
            previousExists: before !== undefined,
            previousContent: before
          }
        ]
      })
    }
  })

  pi.on('tool_result', async (event, ctx) => {
    const toolCallId = toolCallIdFromEvent(event)
    if (!toolCallId) {
      return
    }
    const op = pending.get(toolCallId)
    if (!op) {
      return
    }

    pending.delete(toolCallId)
    if (event.isError) {
      return
    }

    const sessionId = ctx.sessionManager.getSessionId()
    if (!activeSessionId || sessionId !== activeSessionId) {
      reloadForSession(sessionId, ctx.sessionManager)
    }

    const changedFiles = snapshotChangedFiles(op.files)
    if (changedFiles.length === 0) {
      return
    }

    ensureLegacyHistory(state)

    turnLeafId = captureTurnLeafId(ctx.sessionManager, turnLeafId)
    const operation = recordCapturedOperation(
      state,
      {
        ...op,
        files: changedFiles,
        timestamp: Date.now()
      },
      turnLeafId
    )

    pi.appendEntry(EXTENSION_ID, {
      type: 'capture',
      at: Date.now(),
      operationId: operation.id,
      toolName: op.toolName,
      count: changedFiles.length,
      files: changedFiles.map((f) => path.relative(ctx.cwd, f.absPath) || f.absPath)
    })
    if (activeSessionId) {
      saveState(activeSessionId, state)
    }
  })

  const ensureCurrentSessionState = (ctx: {
    sessionManager: {
      getSessionId(): string
      getBranch?: () => SessionEntryLike[]
      getEntries(): SessionEntryLike[]
    }
  }) => {
    const sessionId = ctx.sessionManager.getSessionId()
    if (!activeSessionId || sessionId !== activeSessionId) {
      reloadForSession(sessionId, ctx.sessionManager)
    } else {
      rebuildStackFromHistory(state, ctx.sessionManager)
    }
  }

  const appendUndoEntry = (operationIds: string[], files: string[]) => {
    pi.appendEntry(EXTENSION_ID, {
      type: 'undo',
      at: Date.now(),
      operationIds,
      count: files.length,
      files
    })
  }

  const lastUndoOperationIds = (): string[] => {
    const last = state.history.at(-1)
    return last?.type === 'undo' ? last.operationIds : []
  }

  const undoChanges = (
    ctx: {
      cwd: string
      sessionManager: {
        getSessionId(): string
        getLeafId(): string | null
        getBranch?: () => SessionEntryLike[]
        getEntries(): SessionEntryLike[]
      }
    },
    count: number
  ) => {
    ensureCurrentSessionState(ctx)
    const result = undoRecentOperations(
      state,
      ctx.cwd,
      count,
      currentSessionLeafId(ctx.sessionManager)
    )
    if (result.steps > 0) {
      appendUndoEntry(lastUndoOperationIds(), result.restored)
    }
    if (activeSessionId) {
      saveState(activeSessionId, state)
    }
    return result
  }

  const lastUserMessageTimestamp = (ctx: {
    sessionManager: { getEntries(): SessionEntryLike[]; getBranch?: () => SessionEntryLike[] }
  }): number | null => {
    const entries = ctx.sessionManager.getBranch
      ? ctx.sessionManager.getBranch()
      : ctx.sessionManager.getEntries()
    for (let i = entries.length - 1; i >= 0; i--) {
      const entry = entries[i]
      if (entry?.type !== 'message') {
        continue
      }
      if (entry?.message?.role !== 'user') {
        continue
      }

      const ts = entry?.message?.timestamp
      if (typeof ts === 'number' && Number.isFinite(ts)) {
        return normalizeTimestampMs(ts)
      }

      const parsed = Date.parse(String(entry?.timestamp ?? ''))
      if (Number.isFinite(parsed)) {
        return parsed
      }
    }
    return null
  }

  const getSelectionCandidates = (ctx: {
    sessionManager: {
      getEntries(): SessionEntryLike[]
      getSessionId(): string
      getBranch?: () => SessionEntryLike[]
    }
  }) => {
    ensureCurrentSessionState(ctx)
    if (state.stack.length === 0) {
      return [] as UndoOperation[]
    }

    const lastUserTs = lastUserMessageTimestamp(ctx)
    if (lastUserTs === null) {
      return reverseOperations(state.stack)
    }

    const filtered = state.stack.filter((op) => op.timestamp >= lastUserTs)
    if (filtered.length === 0) {
      return reverseOperations(state.stack)
    }
    return reverseOperations(filtered)
  }

  const undoSelectedChanges = (
    ctx: {
      cwd: string
      sessionManager: {
        getSessionId(): string
        getLeafId(): string | null
        getBranch?: () => SessionEntryLike[]
        getEntries(): SessionEntryLike[]
      }
    },
    selectedIds: Set<string>
  ) => {
    ensureCurrentSessionState(ctx)
    const result = undoSelectedOperations(
      state,
      ctx.cwd,
      selectedIds,
      currentSessionLeafId(ctx.sessionManager)
    )
    if (result.steps > 0) {
      appendUndoEntry(lastUndoOperationIds(), result.restored)
    }
    if (activeSessionId) {
      saveState(activeSessionId, state)
    }
    return result
  }

  registerReverseLastCommands(pi, {
    pending,
    getSelectionCandidates,
    undoChanges,
    undoSelectedChanges
  })
}