import type { ExtensionContext } from "@mariozechner/pi-coding-agent";
import { formatModelRef, sameModelRef } from "../../src/shared/model";
import type { CompactModel, CompactionModelSelection, CompactionState } from "./types";

function modelCost(model: CompactModel): number {
  return model.cost?.output ?? Number.POSITIVE_INFINITY;
}

function apiPenalty(model: CompactModel): number {
  const api = model.api ?? "";
  if (api === "openai-completions") return 0;
  if (api === "anthropic-messages") return 1;
  if (api === "openai-responses") return 2;
  return 3;
}

function compareRemainingModels(left: CompactModel, right: CompactModel): number {
  return (
    apiPenalty(left) - apiPenalty(right) ||
    modelCost(left) - modelCost(right) ||
    (right.maxTokens ?? 0) - (left.maxTokens ?? 0) ||
    (right.contextWindow ?? 0) - (left.contextWindow ?? 0) ||
    formatModelRef(left).localeCompare(formatModelRef(right))
  );
}

function addUnique(target: CompactModel[], seen: Set<string>, model: CompactModel | null | undefined): void {
  if (!model) return;
  const ref = formatModelRef(model).trim();
  if (!ref) return;
  const key = ref.toLowerCase();
  if (seen.has(key)) return;
  seen.add(key);
  target.push(model);
}

function availableCompactionModels(ctx: ExtensionContext): CompactModel[] {
  return (ctx.modelRegistry.getAvailable() as CompactModel[])
    .filter((model) => Array.isArray(model.input) && model.input.includes("text"));
}

function findByReference(models: CompactModel[], reference: string): { model?: CompactModel; error?: string } {
  const trimmed = reference.trim();
  if (!trimmed) return { error: "Empty model reference" };

  const normalized = trimmed.toLowerCase();
  const canonical = models.filter((model) => formatModelRef(model).toLowerCase() === normalized);
  if (canonical.length === 1) {
    return { model: canonical[0] };
  }

  const bare = models.filter((model) => (model.id ?? "").toLowerCase() === normalized);
  if (bare.length === 1) {
    return { model: bare[0] };
  }
  if (bare.length > 1) {
    return { error: `Ambiguous model id \"${trimmed}\". Use provider/model.` };
  }

  return { error: `Model not found: ${trimmed}` };
}

function currentAvailableModel(ctx: ExtensionContext, models: CompactModel[]): CompactModel | null {
  return models.find((model) => sameModelRef(model, ctx.model)) ?? null;
}

export function resolveCompactionModelSelection(ctx: ExtensionContext, state: CompactionState): CompactionModelSelection {
  const available = availableCompactionModels(ctx);
  const availableRefs = available.map((model) => formatModelRef(model));
  const selectedModels: CompactModel[] = [];
  const seen = new Set<string>();
  const warnings: string[] = [];
  const current = currentAvailableModel(ctx, available);
  let source: CompactionModelSelection["source"] = "none";

  const addPreferredModels = (): void => {
    for (const reference of state.config.modelSelection.preferredModels) {
      const resolved = findByReference(available, reference);
      if (resolved.model) {
        if (selectedModels.length === 0) source = "config-preferred";
        addUnique(selectedModels, seen, resolved.model);
      }
    }
  };

  if (state.sessionOverride?.mode === "model") {
    const resolved = findByReference(available, state.sessionOverride.reference);
    if (resolved.model) {
      addUnique(selectedModels, seen, resolved.model);
      source = "session-override";
    } else if (resolved.error) {
      warnings.push(`compaction override ignored: ${resolved.error}`);
    }
  } else if (state.sessionOverride?.mode === "current") {
    if (current) {
      addUnique(selectedModels, seen, current);
      source = "session-override";
    } else {
      warnings.push("compaction override ignored: current model is not available for summarization");
    }
  }

  if (selectedModels.length === 0) {
    if (state.config.modelSelection.mode === "current" && current) {
      addUnique(selectedModels, seen, current);
      source = "current-model";
    }

    if (selectedModels.length === 0) {
      addPreferredModels();
    }

    if (selectedModels.length === 0 && state.config.modelSelection.allowCurrentSessionModel && current) {
      addUnique(selectedModels, seen, current);
      source = "current-model";
    }
  }

  if (state.config.fallback.onFailure === "next-model" || selectedModels.length === 0) {
    addPreferredModels();
    if (state.config.modelSelection.allowCurrentSessionModel && current) {
      addUnique(selectedModels, seen, current);
      if (source === "none") source = "current-model";
    }

    const remaining = available.filter((model) => !seen.has(formatModelRef(model).toLowerCase())).sort(compareRemainingModels);
    for (const model of remaining) {
      addUnique(selectedModels, seen, model);
    }
  }

  if (source === "none" && selectedModels.length > 0) {
    source = "available-fallback";
  }

  return {
    selected: selectedModels[0] ?? null,
    models: selectedModels,
    availableRefs,
    source,
    override: state.sessionOverride,
    warnings,
  };
}

export function resolveModelReferenceInput(ctx: ExtensionContext, reference: string): { model?: CompactModel; reference?: string; error?: string } {
  const resolved = findByReference(availableCompactionModels(ctx), reference);
  if (!resolved.model) return { error: resolved.error ?? `Model not found: ${reference}` };
  return { model: resolved.model, reference: formatModelRef(resolved.model) };
}