import {
  compact,
  type ExtensionAPI,
} from "@earendil-works/pi-coding-agent";
import { loadConfig, parseModelReference } from "./config.js";

function warn(message: string, error?: unknown): void {
  if (error === undefined) {
    console.warn(`[pi-compaction-model] ${message}`);
  } else {
    console.warn(`[pi-compaction-model] ${message}`, error);
  }
}

/**
 * Hook-provided compactions are marked as extension-generated by Pi. Native
 * preparation intentionally does not trust arbitrary extension details, so it
 * will not carry their file lists into a later compaction. Since this extension
 * returns Pi's own details shape, restore those lists before calling native
 * compact() again. Sets make this harmless for native previous entries too.
 */
function restorePreviousFileOperations(
  preparation: {
    fileOps: { read: Set<string>; edited: Set<string> };
  },
  branchEntries: Array<{ type: string; details?: unknown }>,
): void {
  const previous = [...branchEntries].reverse().find((entry) => entry.type === "compaction");
  if (!previous || typeof previous.details !== "object" || previous.details === null) return;

  const details = previous.details as {
    readFiles?: unknown;
    modifiedFiles?: unknown;
  };

  if (Array.isArray(details.readFiles)) {
    for (const path of details.readFiles) {
      if (typeof path === "string") preparation.fileOps.read.add(path);
    }
  }

  if (Array.isArray(details.modifiedFiles)) {
    for (const path of details.modifiedFiles) {
      if (typeof path === "string") preparation.fileOps.edited.add(path);
    }
  }
}

export default function compactionModel(pi: ExtensionAPI): void {
  pi.on("session_before_compact", async (event, ctx) => {
    const config = loadConfig(ctx);
    if (!config || !config.reasons.includes(event.reason)) return;

    const reference = parseModelReference(config.model);
    if (!reference) {
      warn(`Invalid model '${config.model}'; expected provider/model. Using Pi's active model.`);
      return;
    }

    const model = ctx.modelRegistry.find(reference.provider, reference.modelId);
    if (!model) {
      warn(`Model not found: ${config.model}. Using Pi's active model.`);
      return;
    }

    try {
      const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
      if (!auth.ok) {
        warn(`Authentication failed for ${config.model}: ${auth.error}. Using Pi's active model.`);
        return;
      }

      restorePreviousFileOperations(event.preparation, event.branchEntries);

      const result = await compact(
        event.preparation,
        model,
        auth.apiKey,
        auth.headers,
        event.customInstructions,
        event.signal,
        config.thinkingLevel,
        undefined,
        auth.env,
      );

      return { compaction: result };
    } catch (error) {
      if (!event.signal.aborted) {
        warn(`Compaction with ${config.model} failed; using Pi's active model.`, error);
      }
      return;
    }
  });
}

export {
  COMPACTION_REASONS,
  THINKING_LEVELS,
  parseModelReference,
  resolveConfig,
  type CompactionModelConfig,
  type CompactionReason,
  type ThinkingLevel,
} from "./config.js";