import type { ExtensionAPI, ExtensionContext, InputEvent } from '@earendil-works/pi-coding-agent';
import { completionItems, mentionToken } from './autocomplete.js';
import { discoverParticipants, type ParticipantDirectory } from './participants.js';
import { markerFor, MENTION_CUSTOM_TYPE, recipientsOf } from './protocol.js';

const DIRECTORY_TTL_MS = 60_000;

type PendingMention = { recipients: ReturnType<typeof recipientsOf> };

function messageText(entry: unknown): string {
  const message = (entry as { message?: { role?: string; content?: unknown } })?.message;
  if (message?.role !== 'user') return '';
  if (typeof message.content === 'string') return message.content;
  if (!Array.isArray(message.content)) return '';
  return message.content.flatMap((part) => typeof part?.text === 'string' ? [part.text] : []).join('\n');
}

export function createExtension(options: { discover?: typeof discoverParticipants } = {}) {
  const discover = options.discover ?? discoverParticipants;
  return function messaging(pi: ExtensionAPI): void {
    let directory: ParticipantDirectory | null = null;
    let loadedAt = 0;
    let loading: Promise<ParticipantDirectory | null> | null = null;
    const pending: PendingMention[] = [];

    const refresh = async (ctx: ExtensionContext, force = false) => {
      if (!force && directory && Date.now() - loadedAt < DIRECTORY_TTL_MS) return directory;
      if (!loading) {
        loading = discover(ctx.cwd)
          .then((result) => {
            directory = result;
            loadedAt = Date.now();
            return result;
          })
          .finally(() => { loading = null; });
      }
      return loading;
    };

    pi.on('session_start', async (_event, ctx) => {
      try { await refresh(ctx, true); }
      catch (error) { ctx.ui.notify(`Mention discovery unavailable: ${error instanceof Error ? error.message : String(error)}`, 'warning'); }

      ctx.ui.addAutocompleteProvider((current) => ({
        triggerCharacters: ['@'],
        async getSuggestions(lines, cursorLine, cursorCol, request) {
          const token = mentionToken((lines[cursorLine] ?? '').slice(0, cursorCol));
          if (token === null) return current.getSuggestions(lines, cursorLine, cursorCol, request);
          try { await refresh(ctx); } catch { /* retain the last successful directory */ }
          const items = completionItems(directory?.participants ?? [], token);
          return items.length ? { prefix: `@${token}`, items } : current.getSuggestions(lines, cursorLine, cursorCol, request);
        },
        applyCompletion(lines, cursorLine, cursorCol, item, prefix) {
          return current.applyCompletion(lines, cursorLine, cursorCol, item, prefix);
        },
        shouldTriggerFileCompletion(lines, cursorLine, cursorCol) {
          return current.shouldTriggerFileCompletion?.(lines, cursorLine, cursorCol) ?? true;
        },
      }));
    });

    pi.on('input', async (event: InputEvent, ctx) => {
      if (event.source === 'extension') return { action: 'continue' as const };
      try { await refresh(ctx); } catch { return { action: 'continue' as const }; }
      const recipients = recipientsOf(event.text, directory?.participants ?? []);
      if (recipients.length) pending.push({ recipients });
      return { action: 'continue' as const };
    });

    pi.on('before_agent_start', (_event, ctx) => {
      const next = pending.shift();
      if (!next || !directory) return;
      const branch = ctx.sessionManager.getBranch();
      const target = [...branch].reverse().find((entry) => messageText(entry) !== '');
      const persistedRecipients = target ? recipientsOf(messageText(target), directory.participants) : [];
      if (!target || next.recipients.some((recipient) => !persistedRecipients.some((persisted) => persisted.id === recipient.id))) {
        ctx.ui.notify('Mention metadata was not recorded because the submitted message could not be located', 'warning');
        return;
      }
      const duplicate = branch.some((entry) => entry.type === 'custom'
        && entry.customType === MENTION_CUSTOM_TYPE
        && (entry.data as { target?: { entryId?: unknown } } | undefined)?.target?.entryId === target.id);
      if (duplicate) return;
      pi.appendEntry(MENTION_CUSTOM_TYPE, markerFor(directory.actor, next.recipients, ctx.sessionManager.getSessionId(), target.id));
    });
  };
}

export default createExtension();