/// <reference path="./ambient-modules.d.ts" />
import * as fs from 'node:fs'
import * as path from 'node:path'
import { getAgentDir, resolvePathFromAgentDir } from '@firstpick/pi-utils'

export type FileSnapshot = {
  absPath: string
  previousExists: boolean
  previousContent?: string
  afterExists?: boolean
  afterContent?: string
}

export type FileCheckoutAction =
  | { kind: 'write'; absPath: string; content: string }
  | { kind: 'delete'; absPath: string }

type PlannedFileState = { exists: true; content: string } | { exists: false } | { unknown: true }

export type UndoOperation = {
  id: string
  toolName: 'write' | 'edit'
  timestamp: number
  files: FileSnapshot[]
  sessionEntryId?: string
}

export type UndoHistoryRecord =
  | {
      type: 'capture'
      entryId?: string
      operation: UndoOperation
    }
  | {
      type: 'undo'
      entryId?: string
      operationIds: string[]
    }

export type SessionUndoState = {
  stack: UndoOperation[]
  history: UndoHistoryRecord[]
  /** Files the user chose to leave deleted after an outside (non-pi) deletion. */
  pins?: string[]
}

export type SessionEntryLike = {
  id?: string
  type?: string
  timestamp?: string | number
  message?: { role?: string; timestamp?: number }
}

export const MAX_STACK = 100
const MAX_HISTORY = MAX_STACK * 4
const STATE_DIR_NAME = 'rewind-unwind'
const LEGACY_STATE_DIR_NAME = 'reverse-last'

function defaultStateDir(): string {
  return path.join(getAgentDir(), 'state', STATE_DIR_NAME)
}

function legacyStateDir(): string {
  return path.join(getAgentDir(), 'state', LEGACY_STATE_DIR_NAME)
}

function getStateDir(): string {
  const configured = (
    process.env.PI_REWIND_UNWIND_STATE_DIR ?? process.env.PI_REVERSE_LAST_STATE_DIR
  )?.trim()
  if (configured) {
    return path.isAbsolute(configured)
      ? path.normalize(configured)
      : resolvePathFromAgentDir(configured)
  }

  return defaultStateDir()
}

function getStatePath(sessionId: string): string {
  return path.join(getStateDir(), `${sessionId}.json`)
}

function getLegacyStatePath(sessionId: string): string {
  return path.join(legacyStateDir(), `${sessionId}.json`)
}

function parseState(raw: string): SessionUndoState {
  let parsed: SessionUndoState
  try {
    parsed = JSON.parse(raw) as SessionUndoState
  } catch {
    return { stack: [], history: [] }
  }
  if (!parsed || !Array.isArray(parsed.stack)) {
    return { stack: [], history: [] }
  }

  const stack = parsed.stack.filter((item) => item && Array.isArray(item.files)).slice(-MAX_STACK)
  const history = Array.isArray(parsed.history)
    ? parsed.history
        .filter(
          (item) =>
            item &&
            ((item.type === 'capture' && item.operation && Array.isArray(item.operation.files)) ||
              (item.type === 'undo' && Array.isArray(item.operationIds)))
        )
        .slice(-MAX_HISTORY)
    : []
  const pins = Array.isArray(parsed.pins)
    ? parsed.pins.filter((item): item is string => typeof item === 'string')
    : []
  return { stack, history, pins }
}

export function ensureDir(dir: string): void {
  fs.mkdirSync(dir, { recursive: true })
}

export function readTextFile(absPath: string): string | undefined {
  try {
    return fs.readFileSync(absPath, 'utf8')
  } catch {
    return undefined
  }
}

export function resolveToolPath(cwd: string, inputPath: string): string {
  return path.isAbsolute(inputPath) ? path.normalize(inputPath) : path.resolve(cwd, inputPath)
}

export function loadState(sessionId: string): SessionUndoState {
  try {
    return parseState(fs.readFileSync(getStatePath(sessionId), 'utf8'))
  } catch {
    try {
      return parseState(fs.readFileSync(getLegacyStatePath(sessionId), 'utf8'))
    } catch {
      return { stack: [], history: [] }
    }
  }
}

export function saveState(sessionId: string, state: SessionUndoState): void {
  ensureDir(getStateDir())
  const statePath = getStatePath(sessionId)
  const tempPath = `${statePath}.tmp`
  fs.writeFileSync(
    tempPath,
    JSON.stringify({
      stack: state.stack.slice(-MAX_STACK),
      history: state.history.slice(-MAX_HISTORY),
      pins: state.pins ?? []
    }),
    'utf8'
  )
  try {
    fs.renameSync(tempPath, statePath)
  } catch {
    fs.rmSync(statePath, { force: true })
    fs.renameSync(tempPath, statePath)
  }
}

function getBranchEntryIds(sessionManager: {
  getBranch?: () => SessionEntryLike[]
  getEntries(): SessionEntryLike[]
}): Set<string> {
  const entries = sessionManager.getBranch
    ? sessionManager.getBranch()
    : sessionManager.getEntries()
  return new Set(
    entries.map((entry) => entry.id).filter((id): id is string => typeof id === 'string')
  )
}

export function rebuildStackFromHistory(
  state: SessionUndoState,
  sessionManager: { getBranch?: () => SessionEntryLike[]; getEntries(): SessionEntryLike[] }
): void {
  if (state.history.length === 0) {
    return
  }

  const branchEntryIds = getBranchEntryIds(sessionManager)
  const stack: UndoOperation[] = []
  for (const record of state.history) {
    const recordEntryId =
      record.entryId ?? (record.type === 'capture' ? record.operation.sessionEntryId : undefined)
    if (recordEntryId && !branchEntryIds.has(recordEntryId)) {
      continue
    }

    if (record.type === 'capture') {
      stack.push(record.operation)
      continue
    }

    const undone = new Set(record.operationIds)
    for (let index = stack.length - 1; index >= 0; index--) {
      if (undone.has(stack[index]?.id)) {
        stack.splice(index, 1)
      }
    }
  }

  state.stack = stack.slice(-MAX_STACK)
}

function operationsTouchingPath(stack: UndoOperation[], absPath: string): UndoOperation[] {
  return stack.filter((operation) => operation.files.some((file) => file.absPath === absPath))
}

function previousFileState(file: FileSnapshot): PlannedFileState {
  if (file.previousExists) {
    return { exists: true, content: file.previousContent ?? '' }
  }
  return { exists: false }
}

function afterFileState(file: FileSnapshot): PlannedFileState {
  if (file.afterExists === undefined && file.afterContent === undefined) {
    return { unknown: true }
  }
  if (file.afterExists === false) {
    return { exists: false }
  }
  return { exists: true, content: file.afterContent ?? '' }
}

function desiredFileState(
  previousStack: UndoOperation[],
  nextStack: UndoOperation[],
  absPath: string
): PlannedFileState | undefined {
  const nextOps = operationsTouchingPath(nextStack, absPath)
  if (nextOps.length > 0) {
    let state: PlannedFileState | undefined
    for (const operation of nextOps) {
      const file = operation.files.find((item) => item.absPath === absPath)
      if (!file) {
        continue
      }
      if (!state) {
        state = previousFileState(file)
      }
      const after = afterFileState(file)
      if ('unknown' in after) {
        return { unknown: true }
      }
      state = after
    }
    return state
  }

  const previousOps = operationsTouchingPath(previousStack, absPath)
  const first = previousOps[0]?.files.find((item) => item.absPath === absPath)
  return first ? previousFileState(first) : undefined
}

export function planTreeFileCheckout(
  previousStack: UndoOperation[],
  nextStack: UndoOperation[]
): FileCheckoutAction[] {
  const paths = new Set<string>()
  for (const operation of [...previousStack, ...nextStack]) {
    for (const file of operation.files) {
      paths.add(file.absPath)
    }
  }

  const actions: FileCheckoutAction[] = []
  for (const absPath of paths) {
    const desired = desiredFileState(previousStack, nextStack, absPath)
    if (!desired || 'unknown' in desired) {
      continue
    }
    if (desired.exists) {
      actions.push({ kind: 'write', absPath, content: desired.content })
    } else {
      actions.push({ kind: 'delete', absPath })
    }
  }
  return actions
}

function fileMatchesState(absPath: string, state: PlannedFileState): boolean {
  const current = readTextFile(absPath)
  if ('unknown' in state) {
    return true
  }
  if (state.exists) {
    return current === state.content
  }
  return current === undefined
}

function knownFileStates(file: FileSnapshot): PlannedFileState[] {
  const states: PlannedFileState[] = [previousFileState(file)]
  const after = afterFileState(file)
  if (!('unknown' in after)) {
    states.push(after)
  }
  return states
}

export function isUntrackedDiskEdit(
  absPath: string,
  previousStack: UndoOperation[],
  nextStack: UndoOperation[]
): boolean {
  const known = new Map<string, PlannedFileState>()
  for (const operation of [...previousStack, ...nextStack]) {
    const file = operation.files.find((item) => item.absPath === absPath)
    if (!file) {
      continue
    }
    for (const state of knownFileStates(file)) {
      const key = 'exists' in state && state.exists ? `write:${state.content}` : 'delete'
      known.set(key, state)
    }
  }
  if (known.size === 0) {
    return false
  }
  return ![...known.values()].some((state) => fileMatchesState(absPath, state))
}

/** A pin is honored only while the file is still absent from disk. */
export function isPinnedDeleted(pins: string[] | undefined, absPath: string): boolean {
  return (pins ?? []).includes(absPath) && readTextFile(absPath) === undefined
}

/**
 * The file is gone from disk, but a capture still on either branch claims it
 * should exist (afterExists). The plugin itself only deletes files while popping
 * the operation that created them, so this means something outside pi deleted it.
 */
export function isOutsideDeletion(
  absPath: string,
  previousStack: UndoOperation[],
  nextStack: UndoOperation[],
  pins?: string[]
): boolean {
  if (isPinnedDeleted(pins, absPath)) {
    return false
  }
  if (readTextFile(absPath) !== undefined) {
    return false
  }
  return [...previousStack, ...nextStack].some((operation) =>
    operation.files.some((file) => file.absPath === absPath && file.afterExists === true)
  )
}

/** Files deleted outside pi that this jump's plan would silently recreate. */
export function findOutsideDeletions(
  previousStack: UndoOperation[],
  nextStack: UndoOperation[],
  pins?: string[]
): string[] {
  const deletions: string[] = []
  for (const action of planTreeFileCheckout(previousStack, nextStack)) {
    if (action.kind !== 'write') {
      continue
    }
    if (isOutsideDeletion(action.absPath, previousStack, nextStack, pins)) {
      deletions.push(action.absPath)
    }
  }
  return deletions
}

/** Drop pins for files that reappeared on disk (the deletion was undone). */
export function pruneStalePins(state: SessionUndoState): void {
  if (!state.pins) {
    return
  }
  state.pins = state.pins.filter((absPath) => readTextFile(absPath) === undefined)
}

export function findUntrackedDiskEdits(
  previousStack: UndoOperation[],
  nextStack: UndoOperation[],
  pins?: string[]
): string[] {
  const paths = new Set<string>()
  for (const operation of [...previousStack, ...nextStack]) {
    for (const file of operation.files) {
      paths.add(file.absPath)
    }
  }
  return [...paths].filter(
    (absPath) =>
      !isPinnedDeleted(pins, absPath) &&
      !isOutsideDeletion(absPath, previousStack, nextStack, pins) &&
      isUntrackedDiskEdit(absPath, previousStack, nextStack)
  )
}

function adoptDiskIntoSnapshot(file: FileSnapshot, current: string | undefined): void {
  if (current === undefined) {
    file.previousExists = false
    file.previousContent = undefined
    file.afterExists = false
    file.afterContent = undefined
    return
  }
  file.previousExists = true
  file.previousContent = current
  file.afterExists = true
  file.afterContent = current
}

export function adoptUntrackedDiskEdits(operations: UndoOperation[], absPaths: string[]): number {
  const targets = new Set(absPaths)
  let updated = 0
  for (const operation of operations) {
    for (const file of operation.files) {
      if (!targets.has(file.absPath)) {
        continue
      }
      adoptDiskIntoSnapshot(file, readTextFile(file.absPath))
      updated += 1
    }
  }
  return updated
}

export function plannedTreeFileChanges(
  previousStack: UndoOperation[],
  nextStack: UndoOperation[],
  options?: { skipUntracked?: boolean; pins?: string[]; allowOutsideDeletions?: boolean }
): FileCheckoutAction[] {
  const skipUntracked = options?.skipUntracked ?? true
  const allowOutsideDeletions = options?.allowOutsideDeletions ?? false
  return planTreeFileCheckout(previousStack, nextStack).filter((action) => {
    if (isPinnedDeleted(options?.pins, action.absPath)) {
      return false
    }
    const outside = isOutsideDeletion(action.absPath, previousStack, nextStack, options?.pins)
    if (outside && !allowOutsideDeletions) {
      return false
    }
    if (skipUntracked && !outside && isUntrackedDiskEdit(action.absPath, previousStack, nextStack)) {
      return false
    }
    if (action.kind === 'write') {
      return readTextFile(action.absPath) !== action.content
    }
    return fs.existsSync(action.absPath)
  })
}

export function applyTreeFileCheckout(
  previousStack: UndoOperation[],
  nextStack: UndoOperation[],
  cwd: string,
  options?: { skipUntracked?: boolean; pins?: string[]; allowOutsideDeletions?: boolean }
): string[] {
  const restored: string[] = []
  for (const action of plannedTreeFileChanges(previousStack, nextStack, options)) {
    const label = path.relative(cwd, action.absPath) || action.absPath
    if (action.kind === 'write') {
      const current = readTextFile(action.absPath)
      if (current === action.content) {
        continue
      }
      ensureDir(path.dirname(action.absPath))
      fs.writeFileSync(action.absPath, action.content, 'utf8')
      restored.push(label)
      continue
    }

    if (!fs.existsSync(action.absPath)) {
      continue
    }
    fs.unlinkSync(action.absPath)
    restored.push(`${label} (deleted)`)
  }
  return restored
}