import {
  buildExtensionUINotificationMessage,
  buildExtensionUIRequestMessage,
  buildExtensionUISettledMessage,
  EXTENSION_UI_HIGH_FREQUENCY_UPDATE_THROTTLE_MS,
  type ExtensionUIResponsePayload,
  isExtensionUIFireAndForgetMethod,
} from "./extension-ui-contract.js";
import type { ExtensionUIRequestEvent } from "./pi-events.js";
import type { ServerMetricCollector } from "./server-metric-collector.js";
import type { ServerMessage, Session } from "./types.js";

/** Extension UI request from pi SDK or pi-tui mirror bridge. */
export type ExtensionUIRequest = ExtensionUIRequestEvent;

/** Server-side state for a pending first-class ask request. */
export interface PendingAskState {
  requestId: string;
  questionCount: number;
  /** Timestamp when the ask flow was initiated for round-trip timing. */
  initiatedAt: number;
}

interface PendingExtensionUINotificationBroadcast {
  fingerprint: string;
  message: ServerMessage;
  timer: ReturnType<typeof setTimeout>;
}

interface ExtensionUINotificationBroadcastState {
  emittedAt: number;
  fingerprint: string;
  pending?: PendingExtensionUINotificationBroadcast;
}

export interface ExtensionUIState {
  session: Pick<Session, "id">;
  pendingUIRequests: Map<string, ExtensionUIRequest>;
  /** Last persistent extension UI surfaces/status/title, replayed to late focused clients. */
  persistentExtensionUINotifications?: Map<string, ExtensionUIRequest>;
  /** Last high-frequency extension UI notification delivered to focused clients. */
  extensionUINotificationBroadcasts?: Map<string, ExtensionUINotificationBroadcastState>;
  /** Pending first-class ask request awaiting a user response. */
  pendingAsk?: PendingAskState;
}

/** Extension UI response sent to pi. */
export interface ExtensionUIResponse extends ExtensionUIResponsePayload {
  type: "extension_ui_response";
}

export interface ExtensionUIResponseDeliveryOptions {
  deliver: (response: ExtensionUIResponse, request: ExtensionUIRequest) => boolean;
  metrics?: ServerMetricCollector;
  now?: () => number;
  broadcastSettled?: (message: ServerMessage) => void;
}

export function respondToExtensionUIRequest(
  active: ExtensionUIState | undefined,
  response: ExtensionUIResponse,
  options: ExtensionUIResponseDeliveryOptions,
): boolean {
  const req = active?.pendingUIRequests.get(response.id);
  if (!active || !req) {
    return false;
  }

  if (!options.deliver(response, req)) {
    return false;
  }

  settleExtensionUIRequest(active, response.id, {
    cancelled: !!response.cancelled,
    metrics: options.metrics,
    now: options.now,
    broadcastSettled: options.broadcastSettled,
  });
  return true;
}

export function cancelPendingAskRequest(
  active: ExtensionUIState | undefined,
  options: ExtensionUIResponseDeliveryOptions,
): boolean {
  const ask = active?.pendingAsk;
  if (!active || !ask) {
    return false;
  }

  const delivered = respondToExtensionUIRequest(
    active,
    {
      type: "extension_ui_response",
      id: ask.requestId,
      cancelled: true,
    },
    options,
  );
  if (delivered) {
    return true;
  }

  settleExtensionUIRequest(active, ask.requestId, {
    cancelled: true,
    metrics: options.metrics,
    now: options.now,
    broadcastSettled: options.broadcastSettled,
  });
  return false;
}

function notificationReplayKey(req: ExtensionUIRequest): string | undefined {
  switch (req.method) {
    case "setStatus":
      return req.statusKey ? `status:${req.statusKey}` : undefined;
    case "setWidget":
      return req.widgetKey ? `widget:${req.widgetKey}` : undefined;
    case "setTitle":
      return "title";
    case "setWorkingMessage":
      return "working:message";
    case "setWorkingVisible":
      return "working:visible";
    case "setWorkingIndicator":
      return "working:indicator";
    case "setHiddenThinkingLabel":
      return "thinking:hidden-label";
    case "setToolsExpanded":
      return "tools:expanded";
    default:
      return undefined;
  }
}

function widgetLinesHaveContent(lines: string[] | undefined): boolean {
  return lines?.some((line) => line.replace(/[\r\n]/g, "").length > 0) ?? false;
}

function hasPersistentNotificationContent(req: ExtensionUIRequest): boolean {
  switch (req.method) {
    case "setStatus":
      return (req.statusText?.trim().length ?? 0) > 0;
    case "setWidget":
      return req.nativeSurface !== undefined || widgetLinesHaveContent(req.widgetLines);
    case "setTitle":
      return (req.title?.trim().length ?? 0) > 0;
    case "setWorkingMessage":
      return (req.message?.trim().length ?? 0) > 0;
    case "setWorkingVisible":
      return typeof req.workingVisible === "boolean";
    case "setWorkingIndicator":
      return req.workingIndicator !== undefined;
    case "setHiddenThinkingLabel":
      return (req.hiddenThinkingLabel?.trim().length ?? 0) > 0;
    case "setToolsExpanded":
      return typeof req.toolsExpanded === "boolean";
    default:
      return false;
  }
}

export function updatePersistentExtensionUINotifications(
  active: Pick<ExtensionUIState, "persistentExtensionUINotifications">,
  req: ExtensionUIRequest,
): void {
  const key = notificationReplayKey(req);
  if (!key) {
    return;
  }

  const store = (active.persistentExtensionUINotifications ??= new Map());
  if (hasPersistentNotificationContent(req)) {
    if (req.method === "setWidget") {
      store.delete(key);
    }
    store.set(key, req);
    return;
  }

  const clearReq = buildPersistentExtensionUIClearRequest(req);
  if (clearReq) {
    store.set(key, clearReq);
  } else {
    store.delete(key);
  }
}

export function buildPersistentExtensionUINotificationMessages(
  active: Pick<ExtensionUIState, "persistentExtensionUINotifications">,
): ServerMessage[] {
  return Array.from(active.persistentExtensionUINotifications?.values() ?? []).map((req) =>
    buildExtensionUINotificationMessage(req),
  );
}

function buildPersistentExtensionUIClearRequest(
  req: ExtensionUIRequest,
): ExtensionUIRequest | undefined {
  switch (req.method) {
    case "setStatus":
      return req.statusKey
        ? {
            type: "extension_ui_request",
            id: req.id,
            method: "setStatus",
            statusKey: req.statusKey,
            extensionScopeId: req.extensionScopeId,
            extensionDisplayName: req.extensionDisplayName,
          }
        : undefined;
    case "setWidget":
      return req.widgetKey
        ? {
            type: "extension_ui_request",
            id: req.id,
            method: "setWidget",
            widgetKey: req.widgetKey,
            extensionScopeId: req.extensionScopeId,
            extensionDisplayName: req.extensionDisplayName,
          }
        : undefined;
    case "setTitle":
      return { type: "extension_ui_request", id: req.id, method: "setTitle" };
    case "setWorkingMessage":
      return { type: "extension_ui_request", id: req.id, method: "setWorkingMessage" };
    case "setWorkingVisible":
      return {
        type: "extension_ui_request",
        id: req.id,
        method: "setWorkingVisible",
        workingVisible: true,
      };
    case "setWorkingIndicator":
      return { type: "extension_ui_request", id: req.id, method: "setWorkingIndicator" };
    case "setHiddenThinkingLabel":
      return { type: "extension_ui_request", id: req.id, method: "setHiddenThinkingLabel" };
    case "setToolsExpanded":
      return {
        type: "extension_ui_request",
        id: req.id,
        method: "setToolsExpanded",
        toolsExpanded: false,
      };
    default:
      return undefined;
  }
}

export function drainPersistentExtensionUIClearMessages(
  active: Pick<ExtensionUIState, "persistentExtensionUINotifications">,
): ServerMessage[] {
  const store = active.persistentExtensionUINotifications;
  if (!store?.size) {
    return [];
  }

  const messages: ServerMessage[] = [];
  for (const req of store.values()) {
    const clearReq = buildPersistentExtensionUIClearRequest(req);
    if (clearReq) {
      messages.push(buildExtensionUINotificationMessage(clearReq));
    }
  }
  store.clear();
  return messages;
}

function notificationBroadcastThrottleKey(req: ExtensionUIRequest): string | undefined {
  switch (req.method) {
    case "setStatus":
      return req.statusKey ? `status:${req.statusKey}` : undefined;
    case "setWidget":
      return req.widgetKey ? `widget:${req.widgetKey}` : undefined;
    case "setWorkingIndicator":
      return "working:indicator";
    case "setWorkingMessage":
      return "working:message";
    default:
      return undefined;
  }
}

function stableNotificationFingerprint(message: ServerMessage): string {
  return JSON.stringify(message);
}

function cancelPendingExtensionUINotificationBroadcast(
  state: ExtensionUINotificationBroadcastState | undefined,
): void {
  if (!state?.pending) return;
  clearTimeout(state.pending.timer);
  state.pending = undefined;
}

function clearExtensionUINotificationBroadcasts(
  store: Map<string, ExtensionUINotificationBroadcastState> | undefined,
): void {
  for (const state of store?.values() ?? []) {
    cancelPendingExtensionUINotificationBroadcast(state);
  }
  store?.clear();
}

function normalizedNotificationRequestFromMessage(
  req: ExtensionUIRequest,
  message: ServerMessage,
): ExtensionUIRequest {
  if (message.type !== "extension_ui_notification") return req;
  return {
    type: "extension_ui_request",
    id: req.id,
    method: message.method,
    message: message.message,
    notifyType: message.notifyType,
    statusKey: message.statusKey,
    statusText: message.statusText,
    title: message.title,
    text: message.text,
    widgetKey: message.widgetKey,
    widgetLines: message.widgetLines,
    widgetPlacement: message.widgetPlacement,
    extensionScopeId: message.extensionScopeId,
    extensionDisplayName: message.extensionDisplayName,
    workingIndicator: message.workingIndicator,
    workingVisible: message.workingVisible,
    hiddenThinkingLabel: message.hiddenThinkingLabel,
    toolsExpanded: message.toolsExpanded,
    nativeSurface: message.nativeSurface,
  };
}

function isImmediateExtensionUINotification(req: ExtensionUIRequest): boolean {
  if (req.method !== "setWorkingIndicator") return false;
  const indicator = req.workingIndicator;
  if (typeof indicator !== "object" || indicator === null || Array.isArray(indicator)) return false;
  const frames = (indicator as Record<string, unknown>).frames;
  return Array.isArray(frames) && frames.length === 0;
}

function broadcastExtensionUINotification(
  active: ExtensionUIState,
  req: ExtensionUIRequest,
  message: ServerMessage,
  deps: {
    broadcast: (message: ServerMessage) => void;
    now?: () => number;
  },
): void {
  const key = notificationBroadcastThrottleKey(req);
  if (!key) {
    deps.broadcast(message);
    return;
  }

  const store = (active.extensionUINotificationBroadcasts ??= new Map());
  const previous = store.get(key);
  if (!hasPersistentNotificationContent(req)) {
    cancelPendingExtensionUINotificationBroadcast(previous);
    store.delete(key);
    deps.broadcast(message);
    return;
  }

  const now = deps.now?.() ?? Date.now();
  const fingerprint = stableNotificationFingerprint(message);
  if (previous?.fingerprint === fingerprint) {
    cancelPendingExtensionUINotificationBroadcast(previous);
    return;
  }
  if (previous?.pending?.fingerprint === fingerprint) {
    return;
  }

  if (
    !isImmediateExtensionUINotification(req) &&
    previous &&
    now - previous.emittedAt < EXTENSION_UI_HIGH_FREQUENCY_UPDATE_THROTTLE_MS
  ) {
    const delayMs = Math.max(
      0,
      EXTENSION_UI_HIGH_FREQUENCY_UPDATE_THROTTLE_MS - (now - previous.emittedAt),
    );
    cancelPendingExtensionUINotificationBroadcast(previous);
    previous.pending = {
      fingerprint,
      message,
      timer: setTimeout(() => {
        const pending = previous.pending;
        if (!pending) return;
        previous.pending = undefined;
        previous.emittedAt = deps.now?.() ?? Date.now();
        previous.fingerprint = pending.fingerprint;
        deps.broadcast(pending.message);
      }, delayMs),
    };
    return;
  }

  cancelPendingExtensionUINotificationBroadcast(previous);
  store.set(key, { emittedAt: now, fingerprint });
  deps.broadcast(message);
}

export function handleExtensionUIRequest(
  active: ExtensionUIState,
  req: ExtensionUIRequest,
  deps: {
    broadcast: (message: ServerMessage) => void;
    now?: () => number;
  },
): void {
  if (isExtensionUIFireAndForgetMethod(req.method)) {
    const message = buildExtensionUINotificationMessage(req);
    const normalizedReq = normalizedNotificationRequestFromMessage(req, message);
    updatePersistentExtensionUINotifications(active, normalizedReq);
    broadcastExtensionUINotification(active, normalizedReq, message, deps);
    return;
  }

  const broadcastMessage = buildExtensionUIRequestMessage(active.session.id, req);
  active.pendingUIRequests.set(req.id, req);

  if (req.method === "ask") {
    active.pendingAsk = {
      requestId: req.id,
      questionCount: req.questions?.length ?? 0,
      initiatedAt: deps.now?.() ?? Date.now(),
    };
  }

  deps.broadcast(broadcastMessage);
}

function completeAskRequest(
  active: Pick<ExtensionUIState, "pendingAsk" | "session">,
  options: {
    cancelled: boolean;
    metrics?: ServerMetricCollector;
    now?: () => number;
  },
): void {
  const ask = active.pendingAsk;
  if (!ask) {
    return;
  }

  const metrics = options.metrics;
  if (metrics && ask.initiatedAt) {
    metrics.record("server.ask_round_trip_ms", (options.now?.() ?? Date.now()) - ask.initiatedAt, {
      sessionId: active.session.id,
      cancelled: options.cancelled ? "true" : "false",
      questionCount: String(ask.questionCount),
    });
  }

  active.pendingAsk = undefined;
}

export function settleExtensionUIRequest(
  active: ExtensionUIState,
  requestId: string,
  options: {
    cancelled?: boolean;
    metrics?: ServerMetricCollector;
    now?: () => number;
    broadcastSettled?: (message: ServerMessage) => void;
    broadcastIfMissing?: boolean;
  } = {},
): boolean {
  const req = active.pendingUIRequests.get(requestId);
  if (!req) {
    if (options.broadcastIfMissing) {
      options.broadcastSettled?.(buildExtensionUISettledMessage(active.session.id, requestId));
    }
    return false;
  }

  active.pendingUIRequests.delete(requestId);
  if (req.method === "ask" || active.pendingAsk?.requestId === requestId) {
    completeAskRequest(active, {
      cancelled: options.cancelled ?? false,
      metrics: options.metrics,
      now: options.now,
    });
  }

  options.broadcastSettled?.(buildExtensionUISettledMessage(active.session.id, requestId));
  return true;
}

export function settleAllExtensionUIRequests(
  active: ExtensionUIState,
  options: {
    cancelled?: boolean;
    metrics?: ServerMetricCollector;
    now?: () => number;
    broadcastSettled?: (message: ServerMessage) => void;
  } = {},
): void {
  for (const requestId of Array.from(active.pendingUIRequests.keys())) {
    settleExtensionUIRequest(active, requestId, options);
  }
}

export function clearExtensionUIState(
  active: Pick<
    ExtensionUIState,
    | "pendingUIRequests"
    | "persistentExtensionUINotifications"
    | "extensionUINotificationBroadcasts"
    | "pendingAsk"
  >,
): void {
  active.pendingUIRequests.clear();
  active.persistentExtensionUINotifications?.clear();
  clearExtensionUINotificationBroadcasts(active.extensionUINotificationBroadcasts);
  active.pendingAsk = undefined;
}

export function drainExtensionUITeardownMessages(
  active: ExtensionUIState,
  options: {
    cancelled?: boolean;
    metrics?: ServerMetricCollector;
    now?: () => number;
  } = {},
): ServerMessage[] {
  const messages: ServerMessage[] = [];
  settleAllExtensionUIRequests(active, {
    ...options,
    broadcastSettled: (message) => messages.push(message),
  });
  messages.push(...drainPersistentExtensionUIClearMessages(active));
  clearExtensionUINotificationBroadcasts(active.extensionUINotificationBroadcasts);
  return messages;
}

export function buildPendingExtensionUIRequestMessages(
  active: ExtensionUIState | undefined,
): ServerMessage[] {
  if (!active) {
    return [];
  }

  const messages: ServerMessage[] = buildPersistentExtensionUINotificationMessages(active);
  for (const req of active.pendingUIRequests.values()) {
    messages.push(buildExtensionUIRequestMessage(active.session.id, req));
  }
  return messages;
}