import type { Model } from "@earendil-works/pi-ai";
import {
  buildSessionContext,
  type CompactionEntry,
  type ExtensionAPI,
  type ExtensionContext,
  type SessionEntry,
} from "@earendil-works/pi-coding-agent";
import {
  findActiveRemoteCheckpoint,
  hasForeignAssistantTurn,
  messagesAfterCoverage,
  REMOTE_ENTRY_TYPE,
  type RemoteCompactionAttachment,
} from "./checkpoint.ts";
import {
  hasConfiguredOverride,
  loadCompactionConfig,
  resolveConfiguredTarget,
  resolveCurrentTarget,
  type CompactionConfig,
} from "./config.ts";
import {
  isAbortFailure,
  runLocalCompaction,
} from "./local-compaction.ts";
import {
  extractRequestShape,
  isJsonRecord,
  rewriteProviderPayload,
  type ActiveRemoteReplay,
} from "./openai.ts";
import {
  buildToolsPayload,
  callRemoteCompaction,
  messagesToResponseItems,
  normalizeResponseItemsForPrompt,
  remoteIdentityMatches,
  supportsRemoteCompaction,
  thinkingLevelToReasoning,
  toRemoteCompactionData,
  type RemoteCompactionData,
  type ResponseItem,
  type ResponsesReasoningConfig,
  type ResponsesTextConfig,
} from "./remote-compaction.ts";
import {
  waitForRemoteCompletion,
  type RemoteCompletionTimer,
} from "./remote-completion.ts";
import {
  clearPendingFormal,
  clearSessionState,
  createRuntimeState,
  type RuntimeState,
} from "./state.ts";

function reportConfigWarning(state: RuntimeState, ctx: ExtensionContext, warning?: string): void {
  if (!warning || state.configWarningShown) return;
  state.configWarningShown = true;
  const message = `pi-smart-compaction: ${warning} Using the current Pi compaction defaults.`;
  if (ctx.hasUI) ctx.ui.notify(message, "warning");
  else console.error(message);
}

function loadConfig(state: RuntimeState, ctx: ExtensionContext): CompactionConfig {
  const loaded = loadCompactionConfig();
  reportConfigWarning(state, ctx, loaded.warning);
  return loaded.config;
}

function matchesCompactionCoverage(
  entries: SessionEntry[],
  sourceLeafId: string,
  compactionEntry: CompactionEntry,
): boolean {
  const entriesById = new Map(entries.map((entry) => [entry.id, entry]));
  const persistedCompaction = entriesById.get(compactionEntry.id);
  if (
    !persistedCompaction ||
    persistedCompaction.type !== "compaction" ||
    persistedCompaction.parentId !== compactionEntry.parentId
  ) return false;

  const visited = new Set<string>();
  let current: SessionEntry | undefined = persistedCompaction;
  while (current) {
    if (visited.has(current.id)) return false;
    visited.add(current.id);
    if (current.id === sourceLeafId) return current.id !== compactionEntry.id;
    if (current.id !== compactionEntry.id && current.type === "compaction") return false;
    current = current.parentId === null ? undefined : entriesById.get(current.parentId);
  }
  return false;
}

function isLatestFormalCompaction(branch: SessionEntry[], compactionEntryId: string): boolean {
  return [...branch].reverse().find((entry) => entry.type === "compaction")?.id === compactionEntryId;
}

function buildRemoteReplay(
  branch: SessionEntry[],
  entries: SessionEntry[],
  model: Model<any>,
): ActiveRemoteReplay | undefined {
  const checkpoint = findActiveRemoteCheckpoint(branch, entries);
  if (!checkpoint) return undefined;
  if (!remoteIdentityMatches(checkpoint.identity, model)) return undefined;
  const tailMessages = messagesAfterCoverage(branch, checkpoint);
  if (hasForeignAssistantTurn(tailMessages, checkpoint.identity)) return undefined;
  const tailItems = messagesToResponseItems(tailMessages);
  return {
    checkpointId: checkpoint.checkpointId,
    identity: checkpoint.identity,
    explicitHistory: normalizeResponseItemsForPrompt(
      [...checkpoint.data.replacementHistory, ...tailItems],
      model,
    ),
  };
}

async function runRemote(params: {
  pi: ExtensionAPI;
  ctx: ExtensionContext;
  model: Model<any>;
  input: ResponseItem[];
  reasoning?: ResponsesReasoningConfig;
  text?: ResponsesTextConfig;
  keepRecentTokens: number;
  signal?: AbortSignal;
}): Promise<RemoteCompactionData | undefined> {
  let resolveInterrupted!: () => void;
  const interrupted = new Promise<undefined>((resolve) => {
    resolveInterrupted = () => resolve(undefined);
  });
  const onInterrupted = () => resolveInterrupted();
  if (params.signal?.aborted) onInterrupted();
  else params.signal?.addEventListener("abort", onInterrupted, { once: true });

  try {
    if (params.signal?.aborted) return undefined;
    const sessionId = params.ctx.sessionManager.getSessionId();
    const instructions = params.ctx.getSystemPrompt();
    const tools = buildToolsPayload(params.pi.getAllTools(), params.pi.getActiveTools());
    const reasoning = params.reasoning ?? thinkingLevelToReasoning(params.ctx.thinkingLevel);
    // ModelRegistry auth has no AbortSignal parameter. Stop waiting on abort and
    // gate the transport below; once fetch starts, await its real abort-driven settlement.
    const auth = await Promise.race([
      params.ctx.modelRegistry.getApiKeyAndHeaders(params.model),
      interrupted,
    ]);
    if (!auth?.ok || params.signal?.aborted) return undefined;
    const result = await callRemoteCompaction({
      model: params.model,
      apiKey: auth.apiKey,
      headers: auth.headers,
      sessionId,
      input: params.input,
      instructions,
      tools,
      parallelToolCalls: true,
      reasoning,
      text: params.text,
      keepRecentTokens: params.keepRecentTokens,
      signal: params.signal,
    });
    return toRemoteCompactionData(params.model, result);
  } catch {
    return undefined;
  } finally {
    params.signal?.removeEventListener("abort", onInterrupted);
  }
}

export interface SmartCompactionDependencies {
  remoteCompletionTimer?: RemoteCompletionTimer;
}

function registerSmartCompactionExtension(
  pi: ExtensionAPI,
  dependencies: SmartCompactionDependencies,
): void {
  const state = createRuntimeState();

  pi.on("session_start", (_event, ctx) => {
    clearSessionState(state);
    state.configWarningShown = false;
    loadConfig(state, ctx);
  });

  pi.on("session_before_compact", async (event, ctx) => {
    clearPendingFormal(state);
    if (event.signal.aborted) return { cancel: true };

    const branch = event.branchEntries;
    const sourceLeafId = branch.at(-1)?.id;
    const focused = Boolean(event.customInstructions?.trim());
    if (!focused && sourceLeafId && ctx.model && supportsRemoteCompaction(ctx.model)) {
      try {
        const model = ctx.model;
        const replay = buildRemoteReplay(
          branch,
          ctx.sessionManager.getEntries(),
          model,
        );
        const input = replay
          ? replay.explicitHistory
          : normalizeResponseItemsForPrompt(
              messagesToResponseItems(buildSessionContext(branch, sourceLeafId).messages),
              model,
            );
        const requestShape = state.requestShape;
        const controller = new AbortController();
        const abortFromPi = () => controller.abort(event.signal.reason);
        if (event.signal.aborted) abortFromPi();
        else event.signal.addEventListener("abort", abortFromPi, { once: true });
        const promise = runRemote({
          pi,
          ctx,
          model,
          input,
          reasoning: requestShape?.reasoning,
          text: requestShape?.text,
          keepRecentTokens: event.preparation.settings.keepRecentTokens,
          signal: controller.signal,
        }).finally(() => {
          event.signal.removeEventListener("abort", abortFromPi);
        });
        state.pendingFormal = {
          sourceLeafId,
          reason: event.reason,
          willRetry: event.willRetry,
          coverageThroughEntryId: sourceLeafId,
          controller,
          promise,
        };
      } catch {
        clearPendingFormal(state);
      }
    }

    const config = loadConfig(state, ctx);
    if (!hasConfiguredOverride(config)) return undefined;
    try {
      const target = resolveConfiguredTarget(config, ctx);
      const current = resolveCurrentTarget(ctx);
      if (target.model === current.model && target.reasoningEffort === current.reasoningEffort) {
        return undefined;
      }
      const compaction = await runLocalCompaction({
        target,
        preparation: event.preparation,
        ctx,
        customInstructions: event.customInstructions,
        signal: event.signal,
      });
      return { compaction };
    } catch (error) {
      if (isAbortFailure(error, event.signal)) {
        clearPendingFormal(state);
        return { cancel: true };
      }
      return undefined;
    }
  });

  pi.on("session_compact", async (event, ctx) => {
    const pending = state.pendingFormal;
    state.pendingFormal = undefined;
    if (!pending) return;
    if (
      pending.reason !== event.reason ||
      pending.willRetry !== event.willRetry ||
      !matchesCompactionCoverage(
        ctx.sessionManager.getEntries(),
        pending.sourceLeafId,
        event.compactionEntry,
      ) ||
      !isLatestFormalCompaction(ctx.sessionManager.getBranch(), event.compactionEntry.id)
    ) {
      pending.controller.abort();
      return;
    }

    const remote = await waitForRemoteCompletion({
      promise: pending.promise,
      controller: pending.controller,
      timer: dependencies.remoteCompletionTimer,
    });
    if (!remote) return;
    if (!isLatestFormalCompaction(ctx.sessionManager.getBranch(), event.compactionEntry.id)) return;
    const attachment: RemoteCompactionAttachment = {
      ...remote,
      compactionEntryId: event.compactionEntry.id,
      coverageThroughEntryId: pending.coverageThroughEntryId,
    };
    pi.appendEntry(REMOTE_ENTRY_TYPE, attachment);
  });

  pi.on("before_provider_request", (event, ctx) => {
    if (!ctx.model || !isJsonRecord(event.payload)) return undefined;
    if (supportsRemoteCompaction(ctx.model)) {
      const observedShape = extractRequestShape(event.payload);
      state.requestShape = {
        reasoning: observedShape?.reasoning as ResponsesReasoningConfig | undefined,
        text: observedShape?.text,
      };
    }
    const branch = ctx.sessionManager.getBranch();
    const active = buildRemoteReplay(branch, ctx.sessionManager.getEntries(), ctx.model);
    if (!active) return undefined;
    const rewritten = rewriteProviderPayload({
      payload: event.payload,
      model: ctx.model,
      active,
    });
    return rewritten;
  });

  const clearForNavigation = () => {
    clearPendingFormal(state);
    state.requestShape = undefined;
  };
  pi.on("model_select", clearForNavigation);
  pi.on("thinking_level_select", () => {
    state.requestShape = undefined;
  });
  pi.on("session_tree", clearForNavigation);
  pi.on("session_shutdown", () => clearSessionState(state));
}

export function createSmartCompactionExtension(
  dependencies: SmartCompactionDependencies = {},
): (pi: ExtensionAPI) => void {
  return (pi) => registerSmartCompactionExtension(pi, dependencies);
}

export default function smartCompactionExtension(pi: ExtensionAPI): void {
  registerSmartCompactionExtension(pi, {});
}