import {existsSync} from 'node:fs';
import {homedir} from 'node:os';
import {basename, join} from 'node:path';
import process from 'node:process';
import {stripVTControlCharacters} from 'node:util';
import {
  CustomEditor,
  type ExtensionAPI,
  type ExtensionContext,
  type KeybindingsManager,
} from '@earendil-works/pi-coding-agent';
import {
  type EditorTheme,
  truncateToWidth,
  type TUI,
  visibleWidth,
} from '@earendil-works/pi-tui';

type GitStatus = {
  ahead: number;
  behind: number;
  branch: string;
  commit: string;
  conflicted: number;
  action: string;
  pushAhead: number;
  pushBehind: number;
  remoteBranch: string;
  staged: number;
  stashes: number;
  summary: string;
  tag: string;
  unstaged: number;
  untracked: number;
  upstream: string;
};

const compactNumber = new Intl.NumberFormat('en', {
  maximumFractionDigits: 1,
  notation: 'compact',
});

export function parseGitStatus(output: string): GitStatus {
  const status: GitStatus = {
    ahead: 0,
    behind: 0,
    branch: '',
    commit: '',
    conflicted: 0,
    action: '',
    pushAhead: 0,
    pushBehind: 0,
    remoteBranch: '',
    staged: 0,
    stashes: 0,
    summary: '',
    tag: '',
    unstaged: 0,
    untracked: 0,
    upstream: '',
  };

  for (const line of output.split('\n')) {
    if (line.startsWith('# branch.oid ')) {
      status.commit = line.slice(13);
    } else if (line.startsWith('# branch.head ')) {
      status.branch = line.slice(14) === '(detached)' ? '' : line.slice(14);
    } else if (line.startsWith('# branch.upstream ')) {
      status.upstream = line.slice(18);
      status.remoteBranch = status.upstream.slice(status.upstream.indexOf('/') + 1);
    } else if (line.startsWith('# branch.ab ')) {
      const match = /^# branch\.ab \+(?<ahead>\d+) -(?<behind>\d+)$/v.exec(line);
      status.ahead = Number(match?.groups?.ahead ?? 0);
      status.behind = Number(match?.groups?.behind ?? 0);
    } else if (line.startsWith('# stash ')) {
      status.stashes = Number(line.slice(8));
    } else if (line.startsWith('? ')) {
      status.untracked += 1;
    } else if (line.startsWith('u ')) {
      status.conflicted += 1;
    } else if (line.startsWith('1 ') || line.startsWith('2 ')) {
      status.staged += line[2] === '.' ? 0 : 1;
      status.unstaged += line[3] === '.' ? 0 : 1;
    }
  }

  return status;
}

/**
Installs Sloppi's shell-inspired editor, footer, and status widgets. The editor
also routes Up and Down through Atuin while the user is composing a host command,
so shell history stays separate from Pi's ordinary prompt history.
*/
export default function shellUi(pi: ExtensionAPI): void {
  let gitStatus: GitStatus | undefined;
  let sessionCost = 0;
  let sessionCacheRead = 0;
  let sessionCacheWrite = 0;
  let latestCacheReuse: number | undefined;
  let compactions = 0;
  let requestRender = (): void => undefined;

  const refreshStatus = async (ctx: ExtensionContext): Promise<void> => {
    sessionCost = 0;
    sessionCacheRead = 0;
    sessionCacheWrite = 0;
    latestCacheReuse = undefined;
    compactions = 0;

    for (const entry of ctx.sessionManager.getBranch()) {
      if (entry.type === 'message' && entry.message.role === 'assistant') {
        const {usage} = entry.message;
        const requestInput = usage.input + usage.cacheRead + usage.cacheWrite;
        sessionCost += usage.cost.total;
        sessionCacheRead += usage.cacheRead;
        sessionCacheWrite += usage.cacheWrite;
        latestCacheReuse = requestInput === 0 ? undefined : usage.cacheRead / requestInput;
      } else if (entry.type === 'message' && entry.message.role === 'toolResult') {
        sessionCost += entry.message.usage?.cost.total ?? 0;
      } else if (entry.type === 'compaction') {
        compactions += 1;
        sessionCost += entry.usage?.cost.total ?? 0;
      } else if (entry.type === 'branch_summary') {
        sessionCost += entry.usage?.cost.total ?? 0;
      }
    }

    try {
      const result = await pi.exec('git', ['status', '--porcelain=v2', '--branch', '--show-stash'], {
        cwd: ctx.cwd,
        timeout: 2000,
      });

      gitStatus = result.code === 0 ? parseGitStatus(result.stdout) : undefined;
      if (gitStatus !== undefined) {
        const [summary, tag, gitDirectory, pushBranch] = await Promise.all([
          pi.exec('git', ['log', '-1', '--format=%s'], {cwd: ctx.cwd, timeout: 2000}),
          pi.exec('git', ['describe', '--tags', '--exact-match'], {cwd: ctx.cwd, timeout: 2000}),
          pi.exec('git', ['rev-parse', '--absolute-git-dir'], {cwd: ctx.cwd, timeout: 2000}),
          pi.exec('git', ['rev-parse', '--abbrev-ref', '--symbolic-full-name', '@{push}'], {cwd: ctx.cwd, timeout: 2000}),
        ]);
        gitStatus.summary = summary.code === 0 ? summary.stdout.trim() : '';
        gitStatus.tag = tag.code === 0 ? tag.stdout.trim() : '';

        if (gitDirectory.code === 0) {
          const directory = gitDirectory.stdout.trim();
          const actions = [
            ['rebase', ['rebase-merge', 'rebase-apply']],
            ['merge', ['MERGE_HEAD']],
            ['cherry-pick', ['CHERRY_PICK_HEAD']],
            ['revert', ['REVERT_HEAD']],
            ['bisect', ['BISECT_LOG']],
          ] as const;
          gitStatus.action = actions.find(([, paths]) => paths.some(path => existsSync(join(directory, path))))?.[0] ?? '';
        }

        const push = pushBranch.stdout.trim();
        if (pushBranch.code === 0 && push !== gitStatus.upstream) {
          const divergence = await pi.exec('git', ['rev-list', '--left-right', '--count', `HEAD...${push}`], {
            cwd: ctx.cwd,
            timeout: 2000,
          });
          const [ahead = 0, behind = 0] = divergence.stdout.trim().split(/\s+/v).map(Number);
          gitStatus.pushAhead = divergence.code === 0 ? ahead : 0;
          gitStatus.pushBehind = divergence.code === 0 ? behind : 0;
        }
      }
    } catch {
      gitStatus = undefined;
    }

    requestRender();
  };

  pi.on('session_start', async (_event, ctx) => {
    if (ctx.mode !== 'tui') {
      return;
    }

    await refreshStatus(ctx);
    let renderFullscreenTopStatus: ((width: number) => string) | undefined;
    ctx.ui.setFooter((tui, theme, footerData) => {
      requestRender = (): void => {
        tui.requestRender();
      };

      const statusRows = {
        invalidate: () => undefined,
        render(width: number): string[] {
          const renderRow = (left: string, right: string, hasFill = false): string => {
            const fittedRight = truncateToWidth(right, Math.max(0, width - 1), '');
            const rightWidth = visibleWidth(fittedRight);
            const fittedLeft = truncateToWidth(left, Math.max(0, width - rightWidth - (rightWidth > 0 ? 1 : 0)), '');
            const gap = Math.max(0, width - visibleWidth(fittedLeft) - rightWidth);
            const middle = hasFill && gap > 2
              ? ` ${theme.fg('borderMuted', '─'.repeat(gap - 2))} `
              : ' '.repeat(gap);
            return `${fittedLeft}${middle}${fittedRight}`;
          };

          const extensionStatuses = footerData.getExtensionStatuses();
          const sandboxStatus = extensionStatuses.get('sandbox') ?? theme.fg('warning', 'sandbox ?');
          const modeStatus = extensionStatuses.get('ask-mode') ?? theme.fg('dim', 'agent');
          const additionalStatuses = [...extensionStatuses]
            .filter(([id]) => id !== 'sandbox' && id !== 'ask-mode')
            .map(([, status]) => status);

          let git = '';
          if (gitStatus !== undefined) {
            let reference = gitStatus.branch;
            if (reference.length > 32) {
              reference = `${reference.slice(0, 12)}…${reference.slice(-12)}`;
            }

            if (reference.length === 0) {
              reference = gitStatus.tag.length > 0 ? `#${gitStatus.tag}` : `@${gitStatus.commit.slice(0, 8)}`;
            }

            const remote = gitStatus.remoteBranch.length > 0 && gitStatus.remoteBranch !== gitStatus.branch
              ? `${theme.fg('muted', ':')}${theme.fg('success', gitStatus.remoteBranch)}`
              : '';
            const divergence = `${gitStatus.behind > 0 ? `⇣${gitStatus.behind}` : ''}${gitStatus.ahead > 0 ? `⇡${gitStatus.ahead}` : ''}`;
            const pushDivergence = `${gitStatus.pushBehind > 0 ? `⇠${gitStatus.pushBehind}` : ''}${gitStatus.pushAhead > 0 ? `⇢${gitStatus.pushAhead}` : ''}`;
            const parts = [
              `${theme.fg('syntaxKeyword', '')}  ${theme.fg('success', ` ${reference}`)}${remote}`,
              /(?:^|[^\da-z])wip(?:[^\da-z]|$)/iv.test(gitStatus.summary) ? theme.fg('warning', 'wip') : '',
              divergence.length > 0 ? theme.fg('success', divergence) : '',
              pushDivergence.length > 0 ? theme.fg('success', pushDivergence) : '',
              gitStatus.stashes > 0 ? theme.fg('success', `*${gitStatus.stashes}`) : '',
              gitStatus.action.length > 0 ? theme.fg('error', gitStatus.action) : '',
              gitStatus.conflicted > 0 ? theme.fg('error', `~${gitStatus.conflicted}`) : '',
              gitStatus.staged > 0 ? theme.fg('warning', `+${gitStatus.staged}`) : '',
              gitStatus.unstaged > 0 ? theme.fg('warning', `!${gitStatus.unstaged}`) : '',
              gitStatus.untracked > 0 ? theme.fg('syntaxFunction', `?${gitStatus.untracked}`) : '',
            ].filter(Boolean);
            git = parts.join(' ');
          }

          const osIcons: Partial<Record<NodeJS.Platform, string>> = {
            darwin: '',
            linux: '',
            win32: '',
          };
          const osIcon = theme.fg('text', osIcons[process.platform] ?? process.platform);
          const homeDirectory = homedir();
          const cwd = ctx.cwd === homeDirectory || ctx.cwd.startsWith(`${homeDirectory}/`)
            ? `~${ctx.cwd.slice(homeDirectory.length)}`
            : ctx.cwd;

          const usage = ctx.getContextUsage();
          const contextPercent = usage?.percent;
          let contextColor: 'accent' | 'error' | 'warning' | 'syntaxNumber' = 'accent';

          if (typeof contextPercent === 'number' && contextPercent >= 60) {
            contextColor = 'error';
          } else if (typeof contextPercent === 'number' && contextPercent >= 40) {
            contextColor = 'syntaxNumber';
          } else if (typeof contextPercent === 'number' && contextPercent >= 15) {
            contextColor = 'warning';
          }

          const context = usage === undefined
            ? `${theme.fg('toolTitle', '󰍛')} ${theme.bold(theme.fg('warning', 'ctx ?'))}`
            : `${theme.fg('toolTitle', '󰍛')} ${theme.bold(theme.fg(contextColor, usage.percent === null
              ? '?'
              : `${usage.percent.toFixed(1)}%`))} ${theme.fg('muted', `${usage.tokens === null ? '?' : compactNumber.format(usage.tokens)}/${compactNumber.format(usage.contextWindow)}`)}`;
          const model = theme.fg('syntaxType', ctx.model === undefined
            ? '󰧑 no model'
            : `󰧑 ${ctx.model.provider}/${ctx.model.id}`);
          const cost = theme.fg('muted', `$${sessionCost < 1 ? sessionCost.toFixed(3) : sessionCost.toFixed(2)}`);
          const cache = theme.fg('muted', `cache ${latestCacheReuse === undefined ? '—' : `${(latestCacheReuse * 100).toFixed(0)}%`} R${compactNumber.format(sessionCacheRead)} W${compactNumber.format(sessionCacheWrite)}`);
          const thinking = theme.fg('syntaxKeyword', `󰔏 ${ctx.thinkingLevel ?? 'off'}`);
          const pending = ctx.hasPendingMessages() ? theme.fg('warning', '󰅖 queued') : '';
          const workspacePath = `${osIcon} ${theme.fg('mdHeading', `  ${cwd}`)}`;
          const workspace = `${workspacePath} on ${git}`;
          const essentialStatus = [sandboxStatus, modeStatus, pending].filter(Boolean).join('  ');
          const status = [sandboxStatus, modeStatus, ...additionalStatuses, pending].filter(Boolean).join('  ');
          const compaction = theme.fg('muted', `󰆏 ${compactions}`);
          const compactContext = usage === undefined
            ? `${theme.fg('toolTitle', '󰍛')} ${theme.bold(theme.fg('warning', '?'))}`
            : `${theme.fg('toolTitle', '󰍛')} ${theme.bold(theme.fg(contextColor, usage.percent === null ? '?' : `${usage.percent.toFixed(0)}%`))}`;

          if (width < 60) {
            const directoryName = basename(ctx.cwd);
            const directory = ctx.cwd === homeDirectory
              ? '~'
              : (directoryName.length > 0 ? directoryName : ctx.cwd);
            const compactWorkspace = `${osIcon} ${theme.fg('mdHeading', ` ${directory}`)}`;
            return [
              renderRow([compactWorkspace, git].filter(Boolean).join('  '), ''),
              renderRow([essentialStatus, compactContext].filter(Boolean).join('  '), ''),
            ];
          }

          if (width < 80) {
            return [
              renderRow(workspacePath, git),
              renderRow(essentialStatus, `${compactContext}  ${cost}`),
            ];
          }

          return [
            renderRow(workspace, `${model}  ${thinking}`, true),
            renderRow(status, `${context}  ${cache}  ${cost}  ${compaction}`),
          ];
        },
      };
      const topStatus = {
        invalidate() {
          statusRows.invalidate();
        },
        render(width: number): string[] {
          const plainPrefix = truncateToWidth('├─ ', Math.max(0, width - 1), '');
          const line = statusRows.render(Math.max(1, width - visibleWidth(plainPrefix)))[0] ?? '';
          return [`${theme.fg('dim', plainPrefix)}${line}`];
        },
      };
      const bottomStatus = {
        invalidate() {
          statusRows.invalidate();
        },
        render(width: number): string[] {
          const plainPrefix = truncateToWidth('╰─ ', Math.max(0, width - 1), '');
          const line = statusRows.render(Math.max(1, width - visibleWidth(plainPrefix)))[1] ?? '';
          return [`${theme.fg('dim', plainPrefix)}${line}`];
        },
      };
      const unsubscribe = footerData.onBranchChange(() => {
        void refreshStatus(ctx);
      });
      const unsubscribeUserBash = pi.events.on('sloppi:user-bash-end', () => {
        void refreshStatus(ctx);
      });
      /**
       Returns the first status row for use inside fullscreen's mandatory
       three-row editor allocation, avoiding a blank padded row.
       */
      renderFullscreenTopStatus = (width: number): string => topStatus.render(width)[0] ?? '';
      ctx.ui.setWidget('shell-ui-top', widgetTui => ({
        /**
         Keeps the top status as a widget in regular mode, while fullscreen
         renders it inside the editor's mandatory third dock row.
         */
        invalidate() {
          topStatus.invalidate();
        },
        /**
         Avoids duplicating the status line after the fullscreen editor has
         consumed it to satisfy Pi's three-row editor minimum.
         */
        render(width: number): string[] {
          return widgetTui.mode === 'fullscreen' ? [] : topStatus.render(width);
        },
      }), {placement: 'belowEditor'});
      ctx.ui.setWidget('shell-ui-bottom', () => bottomStatus, {placement: 'belowEditor'});

      return {
        dispose() {
          unsubscribe();
          unsubscribeUserBash();
          renderFullscreenTopStatus = undefined;
          ctx.ui.setWidget('shell-ui-top', undefined);
          ctx.ui.setWidget('shell-ui-bottom', undefined);
        },
        invalidate() {
          topStatus.invalidate();
          bottomStatus.invalidate();
        },
        render: () => [],
      };
    });

    class PromptEditor extends CustomEditor {
      private readonly renderer: TUI;
      private readonly editorKeybindings: KeybindingsManager;
      private shellHistory: {
        commands: string[];
        draft: string;
        index: number;
        prefix: '!' | '!!';
      } | undefined;

      private shellHistoryRequest = 0;

      private atuinWarningShown = false;

      /**
       Retains Pi's editor behavior while keeping access to the live renderer
       proxy so mode switches immediately use the matching layout. The keybinding
       manager is retained so shell history follows remapped Up and Down actions.
       */
      constructor(tui: TUI, theme: EditorTheme, keybindings: KeybindingsManager) {
        super(tui, theme, keybindings, {paddingX: 0});
        this.renderer = tui;
        this.editorKeybindings = keybindings;
      }

      /**
       Searches Atuin when Up is pressed in a `!` or `!!` command and then keeps
       navigation local to the returned snapshot. A request number prevents a
       slow search from replacing text the user edited while it was running.
       All other keys remain owned by CustomEditor so Pi shortcuts, completion,
       paste handling, and ordinary prompt history continue to work unchanged.
       Atuin limits its default CLI query to the newest matches but prints them
       oldest-first; reverse those rows locally rather than using --reverse,
       which changes the database selection to the oldest matches before LIMIT.
       User-only filtering keeps agent tool calls from displacing host inputs.
       */
      handleInput(data: string): void {
        const isUp = this.editorKeybindings.matches(data, 'tui.editor.cursorUp')
          || this.editorKeybindings.matches(data, 'tui.editor.historyPrevious');
        const isDown = this.editorKeybindings.matches(data, 'tui.editor.cursorDown')
          || this.editorKeybindings.matches(data, 'tui.editor.historyNext');
        const currentText = this.getText();
        const currentHistoryText = this.shellHistory === undefined
          ? undefined
          : (this.shellHistory.index === -1
            ? this.shellHistory.draft
            : `${this.shellHistory.prefix}${this.shellHistory.commands[this.shellHistory.index] ?? ''}`);

        if (currentText !== currentHistoryText && this.shellHistory !== undefined) {
          this.shellHistory = undefined;
          this.shellHistoryRequest += 1;
        }

        if ((isUp || isDown) && this.shellHistory !== undefined) {
          const nextIndex = this.shellHistory.index + (isUp ? 1 : -1);
          if (nextIndex >= this.shellHistory.commands.length) {
            return;
          }

          if (nextIndex < 0) {
            this.shellHistory.index = -1;
            this.setText(this.shellHistory.draft);
          } else {
            this.shellHistory.index = nextIndex;
            this.setText(`${this.shellHistory.prefix}${this.shellHistory.commands[nextIndex] ?? ''}`);
          }

          this.renderer.requestRender();
          return;
        }

        const match = /^(?<prefix>!!|!)(?<query>[^\n]*)$/v.exec(currentText);
        if (isUp && match?.groups !== undefined) {
          const prefix: '!' | '!!' = currentText.startsWith('!!') ? '!!' : '!';
          const query = match.groups.query ?? '';
          const request = ++this.shellHistoryRequest;
          const arguments_ = [
            'search',
            '--cmd-only',
            '--print0',
            '--author',
            '$all-user',
            '--limit',
            '100',
            '--search-mode',
            'prefix',
            ...query.length > 0 ? ['--', query] : [],
          ];

          void pi.exec('atuin', arguments_, {cwd: ctx.cwd, timeout: 2000}).then(result => {
            if (request !== this.shellHistoryRequest || this.getText() !== currentText) {
              return;
            }

            if (result.code !== 0) {
              if (result.code !== 1 && !this.atuinWarningShown) {
                this.atuinWarningShown = true;
                ctx.ui.notify('Atuin history is unavailable.', 'warning');
              }

              return;
            }

            const commands = result.stdout.split('\0').filter(command => command.length > 0).toReversed();
            if (commands.length === 0) {
              return;
            }

            this.shellHistory = {
              commands,
              draft: currentText,
              index: 0,
              prefix,
            };
            this.setText(`${prefix}${commands[0] ?? ''}`);
            this.renderer.requestRender();
          }).catch(() => {
            if (request !== this.shellHistoryRequest || this.atuinWarningShown) {
              return;
            }

            this.atuinWarningShown = true;
            ctx.ui.notify('Atuin history is unavailable.', 'warning');
          });
          return;
        }

        if (!isUp && !isDown) {
          this.shellHistory = undefined;
          this.shellHistoryRequest += 1;
        }

        super.handleInput(data);
      }

      /**
       Reframes Pi's bordered editor as the shell prompt. Fullscreen reserves
       at least three editor rows, so its third row carries the top status
       instead of becoming blank; regular mode keeps that status as a widget.
       */
      render(width: number): string[] {
        const plainPrefix = truncateToWidth('╭─❯ ', Math.max(0, width - 1), '');
        const prefixWidth = visibleWidth(plainPrefix);
        const indent = ' '.repeat(prefixWidth);
        const continuationPrefix = prefixWidth === 0 ? '' : `│${' '.repeat(prefixWidth - 1)}`;
        const lines = super.render(Math.max(1, width - prefixWidth));
        const isBorder = (line: string): boolean => /^(?:─+|─── [↑↓] \d+ more ─*)$/v.test(stripVTControlCharacters(line));
        const bottomBorderIndex = lines.findIndex((line, index) => index > 0 && isBorder(line));
        if (bottomBorderIndex === -1) {
          return lines;
        }

        const topBorder = lines[0] ?? '';
        const bottomBorder = lines[bottomBorderIndex] ?? '';
        const prefixCharacters = [...plainPrefix];
        const corner = prefixCharacters.slice(0, 2).join('');
        const arrow = prefixCharacters.slice(2).join('');
        const input = lines.slice(1, bottomBorderIndex).map((line, index) => index === 0
          ? `${ctx.ui.theme.fg('dim', corner)}${this.borderColor(arrow)}${line}`
          : `${ctx.ui.theme.fg('dim', continuationPrefix)}${line}`);
        const autocomplete = lines.slice(bottomBorderIndex + 1).map(line => `${ctx.ui.theme.fg('dim', continuationPrefix)}${line}`);

        return [
          `${this.borderColor('─'.repeat(prefixWidth))}${topBorder}`,
          ...input,
          ...stripVTControlCharacters(bottomBorder).includes('↓') ? [`${indent}${bottomBorder}`] : [],
          ...autocomplete,
          ...renderFullscreenTopStatus !== undefined && this.renderer.mode === 'fullscreen'
            ? [renderFullscreenTopStatus(width)]
            : [],
        ];
      }
    }

    ctx.ui.setEditorComponent((tui, theme, keybindings) => new PromptEditor(tui, theme, keybindings));
  });

  pi.on('turn_end', async (_event, ctx) => {
    if (ctx.mode === 'tui') {
      await refreshStatus(ctx);
    }
  });

  pi.on('session_tree', async (_event, ctx) => {
    if (ctx.mode === 'tui') {
      await refreshStatus(ctx);
    }
  });

  pi.on('session_compact', async (_event, ctx) => {
    if (ctx.mode === 'tui') {
      await refreshStatus(ctx);
    }
  });
}