import {
  resolveInitialModel,
  resolveModelForAgent,
} from "@/lib/chat/use-chat-preferences";
import type { LlmModel } from "@/lib/llm-models.query";
import type { SupportedProvider } from "@/lib/llm-provider-api-keys.query";

type AgentInfo = {
  id: string;
  modelId?: string | null;
  llmApiKeyId?: string | null;
};

type ChatApiKeyInfo = {
  id: string;
  provider: string;
};

type OrganizationInfo = {
  defaultModelId?: string | null;
  defaultLlmApiKeyId?: string | null;
} | null;

/** The current user's saved default (model, key) pair — the "member" level. */
type MemberDefaultInfo = {
  modelId?: string | null;
  chatApiKeyId?: string | null;
} | null;

/** A model identifier is the models.id UUID throughout the chat model flow. */
export type ResolvedInitialAgentState = {
  agentId: string;
  modelId: string;
  apiKeyId: string | null;
};

export type ResolvedChatModelState = {
  modelId: string;
  apiKeyId: string | null;
  provider: SupportedProvider | undefined;
};

export type CreateConversationInput = {
  agentId: string;
  modelId?: string;
  chatApiKeyId?: string | null;
};

export function resolveInitialAgentSelection<TAgent extends AgentInfo>(params: {
  agents: TAgent[];
  organizationDefaultAgentId?: string | null;
  savedAgentId?: string | null;
  memberDefaultAgentId?: string | null;
  canUseSavedAgent: boolean;
}): TAgent | null {
  const { agents } = params;
  if (agents.length === 0) {
    return null;
  }

  const organizationDefaultAgent = agents.find(
    (agent) => agent.id === params.organizationDefaultAgentId,
  );
  if (organizationDefaultAgent) {
    return organizationDefaultAgent;
  }

  if (params.canUseSavedAgent) {
    const savedAgent = agents.find((agent) => agent.id === params.savedAgentId);
    if (savedAgent) {
      return savedAgent;
    }
  }

  const memberDefaultAgent = agents.find(
    (agent) => agent.id === params.memberDefaultAgentId,
  );
  if (memberDefaultAgent) {
    return memberDefaultAgent;
  }

  return agents[0];
}

export function resolveInitialAgentState(params: {
  agent: AgentInfo;
  modelsByProvider: Record<string, LlmModel[]>;
  chatApiKeys: ChatApiKeyInfo[];
  organization: OrganizationInfo;
  memberDefault: MemberDefaultInfo;
}): ResolvedInitialAgentState | null {
  const resolved = resolveChatModelState({
    agent: params.agent,
    modelsByProvider: params.modelsByProvider,
    chatApiKeys: params.chatApiKeys,
    organization: params.organization,
    memberDefault: params.memberDefault,
  });

  if (!resolved) {
    return null;
  }

  return {
    agentId: params.agent.id,
    modelId: resolved.modelId,
    apiKeyId: resolved.apiKeyId,
  };
}

/** Resolve the provider for a model UUID. */
export function getProviderForModelId(params: {
  modelId: string;
  chatModels: LlmModel[];
}): SupportedProvider | undefined {
  return params.chatModels.find((model) => model.dbId === params.modelId)
    ?.provider;
}

export function resolveChatModelState(params: {
  agent: AgentInfo | null;
  modelsByProvider: Record<string, LlmModel[]>;
  chatApiKeys: ChatApiKeyInfo[];
  organization: OrganizationInfo;
  memberDefault: MemberDefaultInfo;
  chatModels?: LlmModel[];
}): ResolvedChatModelState | null {
  // The resolver identifies models by their models.id UUID.
  const modelsByProvider = Object.fromEntries(
    Object.entries(params.modelsByProvider).map(([provider, models]) => [
      provider,
      models.map((m) => ({ id: m.dbId, isBest: m.isBest })),
    ]),
  );

  const resolved = params.agent
    ? resolveModelForAgent({
        agent: params.agent,
        context: {
          modelsByProvider,
          chatApiKeys: params.chatApiKeys,
          organization: params.organization,
          memberDefault: params.memberDefault,
        },
      })
    : resolveInitialModel({
        modelsByProvider,
        chatApiKeys: params.chatApiKeys,
        organization: params.organization,
        memberDefault: params.memberDefault,
        agent: null,
      });

  if (!resolved) {
    return null;
  }

  return {
    modelId: resolved.modelId,
    apiKeyId: resolved.apiKeyId,
    provider:
      params.chatModels && params.chatModels.length > 0
        ? getProviderForModelId({
            modelId: resolved.modelId,
            chatModels: params.chatModels,
          })
        : undefined,
  };
}

export function resolvePreferredModelForProvider(params: {
  provider: SupportedProvider;
  modelsByProvider: Record<string, LlmModel[]>;
}): { modelId: string; provider: SupportedProvider } | null {
  const providerModels = params.modelsByProvider[params.provider];
  if (!providerModels || providerModels.length === 0) {
    return null;
  }

  const bestModel = providerModels.find((model) => model.isBest);

  return {
    modelId: bestModel?.dbId ?? providerModels[0].dbId,
    provider: params.provider,
  };
}

export function buildCreateConversationInput(params: {
  agentId: string | null;
  modelId: string;
  chatApiKeyId: string | null;
}): CreateConversationInput | null {
  if (!params.agentId) {
    return null;
  }

  return {
    agentId: params.agentId,
    modelId: params.modelId || undefined,
    chatApiKeyId: params.chatApiKeyId ?? undefined,
  };
}

export function shouldResetInitialChatState(params: {
  previousRouteConversationId?: string;
  routeConversationId?: string;
}): boolean {
  return !params.routeConversationId && !!params.previousRouteConversationId;
}
