import type { SupportedProvider } from "@shared";
import logger from "@/logging";
import {
  AgentTeamModel,
  ModelModel,
  OptimizationRuleModel,
  TeamModel,
} from "@/models";
import {
  getTokenizer,
  type ProviderMessage,
  type Tokenizer,
} from "@/tokenizers";
import type {
  Agent,
  Anthropic,
  Cerebras,
  Cohere,
  CommonMcpToolDefinition,
  DeepSeek,
  Gemini,
  Groq,
  Minimax,
  Mistral,
  OpenAi,
  Openrouter,
  Perplexity,
  Vllm,
  Xai,
  Zhipuai,
} from "@/types";

type ProviderMessages = {
  anthropic: Anthropic.Types.MessagesRequest["messages"];
  cerebras: Cerebras.Types.ChatCompletionsRequest["messages"];
  cohere: Cohere.Types.ChatRequest["messages"];
  gemini: Gemini.Types.GenerateContentRequest["contents"];
  groq: Groq.Types.ChatCompletionsRequest["messages"];
  openrouter: Openrouter.Types.ChatCompletionsRequest["messages"];
  mistral: Mistral.Types.ChatCompletionsRequest["messages"];
  perplexity: Perplexity.Types.ChatCompletionsRequest["messages"];
  minimax: Minimax.Types.ChatCompletionsRequest["messages"];
  openai: OpenAi.Types.ChatCompletionsRequest["messages"];
  vllm: Vllm.Types.ChatCompletionsRequest["messages"];
  ollama: Vllm.Types.ChatCompletionsRequest["messages"];
  xai: Xai.Types.ChatCompletionsRequest["messages"];
  zhipuai: Zhipuai.Types.ChatCompletionsRequest["messages"];
  deepseek: DeepSeek.Types.ChatCompletionsRequest["messages"];
};

/**
 * Estimate token count for tool definitions by serializing them
 * and using the provider-specific tokenizer for accurate counting.
 */
export function estimateToolTokens(
  tools: CommonMcpToolDefinition[],
  tokenizer: Tokenizer,
): number {
  if (tools.length === 0) return 0;
  const serialized = tools
    .map((t) => {
      let text = t.name;
      if (t.description) text += ` ${t.description}`;
      if (t.inputSchema) text += ` ${JSON.stringify(t.inputSchema)}`;
      return text;
    })
    .join(" ");
  return tokenizer.countTokens({
    role: "user",
    content: serialized,
  } as ProviderMessage);
}

/**
 * Get optimized model based on dynamic optimization rules
 * Returns the optimized model name or null if no optimization applies
 */
export async function getOptimizedModel<
  Provider extends keyof ProviderMessages,
>(
  agent: Agent,
  messages: ProviderMessages[Provider],
  provider: Provider,
  hasTools: boolean,
  tools: CommonMcpToolDefinition[] = [],
): Promise<string | null> {
  const agentId = agent.id;

  // Get organizationId the same way limits do: from agent's teams OR fallback
  let organizationId: string | null = null;
  const agentTeamIds = await AgentTeamModel.getTeamsForAgent(agentId);

  if (agentTeamIds.length > 0) {
    // Get organizationId from agent's first team
    const teams = await TeamModel.findByIds(agentTeamIds);
    if (teams.length > 0 && teams[0].organizationId) {
      organizationId = teams[0].organizationId;
      logger.info(
        { agentId, organizationId },
        "[CostOptimization] resolved organizationId from team",
      );
    }
  } else {
    // If agent has no teams, check if there are any organization optimization rules to apply (fallback)
    // TODO: this fallback doesn't work if there are multiple organizations.
    organizationId = await OptimizationRuleModel.getFirstOrganizationId();

    if (organizationId) {
      logger.info(
        { agentId, organizationId },
        "[CostOptimization] agent has no teams - using fallback organization",
      );
    }
  }

  if (!organizationId) {
    logger.warn(
      { agentId },
      "[CostOptimization] could not resolve organizationId",
    );
    return null;
  }

  // Fetch enabled optimization rules for this organization, agent, and provider
  const rules =
    await OptimizationRuleModel.findEnabledByOrganizationAndProvider(
      organizationId,
      provider,
    );

  if (rules.length === 0) {
    logger.info(
      { agentId, organizationId, provider },
      "[CostOptimization] no optimization rules configured",
    );
    return null;
  }

  // Use provider-specific tokenizer to count tokens
  const tokenizer = getTokenizer(provider);
  const messageTokenCount = tokenizer.countTokens(messages);
  const toolTokenCount = estimateToolTokens(tools, tokenizer);
  const tokenCount = messageTokenCount + toolTokenCount;

  logger.info(
    { tokenCount, messageTokenCount, toolTokenCount, hasTools },
    "[CostOptimization] LLM request evaluated",
  );

  // Evaluate rules and return optimized model (or null if no rule matches)
  const optimizedModel = OptimizationRuleModel.matchByRules(rules, {
    tokenCount,
    hasTools,
  });

  if (optimizedModel) {
    logger.info(
      { agentId, optimizedModel },
      "[CostOptimization] optimization rule matched",
    );
  } else {
    logger.info({ agentId }, "[CostOptimization] no optimization rule matched");
  }

  return optimizedModel;
}

/**
 * Calculate cost for token usage based on model pricing.
 * Uses provider to disambiguate models with the same name across providers.
 * Returns undefined if token counts are not available.
 */
export async function calculateCost(
  model: string,
  inputTokens: number | null | undefined,
  outputTokens: number | null | undefined,
  provider: SupportedProvider,
): Promise<number | undefined> {
  if (!inputTokens || !outputTokens) {
    return undefined;
  }

  const model_entry = await ModelModel.findByProviderAndModelId(
    provider,
    model,
  );
  const pricing = ModelModel.getEffectivePricing(model_entry, model);

  const inputCost =
    (inputTokens / 1_000_000) * Number.parseFloat(pricing.pricePerMillionInput);
  const outputCost =
    (outputTokens / 1_000_000) *
    Number.parseFloat(pricing.pricePerMillionOutput);

  return inputCost + outputCost;
}
