import type { ExtensionContext, SessionBeforeCompactEvent } from "@mariozechner/pi-coding-agent";
import { errorMessage, isAbortError } from "../../src/shared/errors";
import { formatModelRef } from "../../src/shared/model";
import { notify } from "../../src/shared/ui";
import { resolveCompactionModelSelection } from "./model-selection";
import { chunkMessagesForSummary, extractMessagesFromEntries, normalizeMessagesForSummary } from "./prepare";
import { generateStructuredSummaryInChunks } from "./summary";
import { COMPACTION_DETAILS_VERSION, type CompactionRunInfo, type CompactionState } from "./types";

function createFailureRunInfo(state: CompactionState, message: string): CompactionRunInfo | null {
  if (!state.lastRun) return null;
  return { ...state.lastRun, lastError: message };
}

export async function runCompaction(
  event: SessionBeforeCompactEvent,
  ctx: ExtensionContext,
  state: CompactionState,
): Promise<any> {
  if (!state.config.enabled) return;

  const { preparation, branchEntries, signal } = event;
  let messages = state.config.strategy.usePiPreparationFirst
    ? [...preparation.messagesToSummarize, ...preparation.turnPrefixMessages]
    : [];
  let firstKeptEntryId = preparation.firstKeptEntryId;
  let recoveredFromBranchEntries = false;
  let groupedTurns = 0;
  let semanticUnits = 0;

  if (messages.length === 0 && state.config.strategy.recoverFromBranchEntriesWhenEmpty && branchEntries.length > 0) {
    const extracted = extractMessagesFromEntries(branchEntries, {
      keptBudgetChars: state.config.strategy.keptBudgetChars,
      groupByTurns: state.config.strategy.groupByTurns,
      groupToolCycles: state.config.strategy.groupToolCycles,
    });

    groupedTurns = extracted.groupedTurns;
    semanticUnits = extracted.semanticUnits;
    if (extracted.messages.length > 0) {
      messages = extracted.messages;
      firstKeptEntryId = extracted.firstKeptEntryId;
      recoveredFromBranchEntries = true;
    }
  }

  if (messages.length === 0) {
    notify(ctx, "compaction skipped: session is too small to summarize safely", "info");
    return { cancel: true };
  }

  const normalized = normalizeMessagesForSummary(messages, state.config.strategy);
  normalized.stats.recoveredFromBranchEntries = recoveredFromBranchEntries;
  normalized.stats.groupedTurns = groupedTurns;
  if (semanticUnits > 0) normalized.stats.semanticUnits = semanticUnits;

  if (normalized.messages.length === 0) {
    notify(ctx, "compaction skipped: no usable messages remained after normalization", "info");
    return { cancel: true };
  }

  const selection = resolveCompactionModelSelection(ctx, state);
  for (const warning of selection.warnings) {
    notify(ctx, warning, "warning");
  }

  if (!selection.selected || selection.models.length === 0) {
    notify(ctx, "compaction failed: no configured model is available for summarization", "error");
    state.lastRun = createFailureRunInfo(state, "No configured compaction model available");
    return { cancel: true };
  }

  const attempts = selection.models.slice(0, Math.max(1, state.config.fallback.maxAttempts));
  const failures: string[] = [];

  for (let index = 0; index < attempts.length; index++) {
    const model = attempts[index];
    const modelRef = formatModelRef(model);
    const chunks = chunkMessagesForSummary(normalized.messages, state.config.strategy, normalized.stats);

    if (state.config.notifications.showModelUsed) {
      const chunkSuffix = chunks.length > 1 ? `, ${chunks.length} chunks` : "";
      notify(
        ctx,
        `compacting ${normalized.messages.length} messages with ${modelRef} (${index + 1}/${attempts.length}${chunkSuffix})`,
        "info",
      );
    }

    const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
    if (!auth.ok) {
      failures.push(`${modelRef}: ${auth.error}`);
      continue;
    }

    try {
      const summaryResult = await generateStructuredSummaryInChunks(
        chunks,
        model,
        preparation.settings.reserveTokens,
        { apiKey: auth.apiKey, headers: auth.headers },
        signal,
        event.customInstructions,
        preparation.previousSummary,
      );
      const summary = summaryResult.summary;

      if (!summary?.trim()) {
        failures.push(`${modelRef}: empty summary`);
        continue;
      }

      normalized.stats.chunkCount = summaryResult.chunkCount;
      state.lastRun = {
        modelReference: modelRef,
        source: selection.source,
        attempt: index + 1,
        attempts: attempts.length,
        usedFallback: index > 0,
        compression: normalized.stats,
      };

      if (state.config.notifications.showCompressionStats) {
        notify(
          ctx,
          `compaction stats: ${normalized.stats.outputMessages}/${normalized.stats.inputMessages} msgs · units ${normalized.stats.semanticUnits} · chunks ${normalized.stats.chunkCount} · thinking -${normalized.stats.droppedThinkingBlocks} · toolResult compressed ${normalized.stats.compressedToolResults} · snip ${normalized.stats.snippedSemanticUnits}u/${normalized.stats.snippedMessages}m · micro ${normalized.stats.microcompactedToolResults} · collapse ${normalized.stats.collapsedSemanticUnits}u/${normalized.stats.collapsedMessages}m`,
          "info",
        );
      }

      return {
        compaction: {
          summary,
          firstKeptEntryId,
          tokensBefore: preparation.tokensBefore,
          details: {
            version: COMPACTION_DETAILS_VERSION,
            model: {
              provider: model.provider,
              id: model.id,
              reference: modelRef,
            },
            selection: {
              source: selection.source,
              override: selection.override,
              candidates: attempts.map((candidate) => formatModelRef(candidate)),
            },
            input: {
              source: recoveredFromBranchEntries ? "branch-recovery" : "pi-preparation",
              groupedTurns,
              semanticUnits: normalized.stats.semanticUnits,
            },
            compression: normalized.stats,
          },
        },
      };
    } catch (error) {
      if (signal.aborted || isAbortError(error)) {
        return { cancel: true };
      }
      failures.push(`${modelRef}: ${errorMessage(error)}`);
      if (state.config.fallback.onFailure !== "next-model") {
        break;
      }
    }
  }

  const failureText = failures.join(" | ") || "unknown failure";
  state.lastRun = {
    modelReference: formatModelRef(selection.selected),
    source: selection.source,
    attempt: attempts.length,
    attempts: attempts.length,
    usedFallback: attempts.length > 1,
    compression: normalized.stats,
    lastError: failureText,
  };
  notify(ctx, `compaction failed: ${failureText}`, "error");
  return { cancel: true };
}