import type { AIProviderRegistryStatusEntry, ChatMessage } from '@genesislcap/foundation-ai';
import type { CostSessionModelEntry } from './cost-session-history';

const mergeModelEntry = (existing: CostSessionModelEntry, message: ChatMessage): void => {
  if (existing.provider === undefined && message.provider !== undefined) {
    existing.provider = message.provider;
  }
};

const collectFromMessages = (
  messages: readonly ChatMessage[],
  seen: Map<string, CostSessionModelEntry>,
): void => {
  for (const m of messages) {
    if (m.role === 'assistant' && m.model) {
      const existing = seen.get(m.model);
      if (!existing) {
        seen.set(m.model, {
          model: m.model,
          provider: m.provider,
        });
      } else {
        mergeModelEntry(existing, m);
      }
    }
    for (const tc of m.toolCalls ?? []) {
      if (tc.subAgentTrace) collectFromMessages(tc.subAgentTrace, seen);
    }
  }
};

const providerByModelFromStatuses = (
  providerStatuses: readonly AIProviderRegistryStatusEntry[],
): Map<string, NonNullable<CostSessionModelEntry['provider']>> => {
  const map = new Map<string, NonNullable<CostSessionModelEntry['provider']>>();
  for (const entry of providerStatuses) {
    const model = entry.status?.model;
    const provider = entry.status?.provider;
    if (model && provider) {
      map.set(model, provider);
    }
  }
  return map;
};

/** Unique models used across a message list (includes sub-agent traces). */
export function collectSessionModels(
  messages: readonly ChatMessage[],
  providerStatuses: readonly AIProviderRegistryStatusEntry[] = [],
): CostSessionModelEntry[] {
  if (providerStatuses.length === 0) {
    console.warn(
      'collectSessionModels called without providerStatuses. Registry provider fallback unavailable.',
    );
  }

  const seen = new Map<string, CostSessionModelEntry>();
  collectFromMessages(messages, seen);
  const providerByModel = providerByModelFromStatuses(providerStatuses);

  return [...seen.values()].map(({ model, provider }) => ({
    model,
    provider: provider ?? providerByModel.get(model),
  }));
}
