import { Injectable } from "@nestjs/common";
import { OpentelemetryProvider } from "../opentelemetry/opentelemetry.provider";
import {
  MODEL_COST_PER_1K_COMPLETION_TOKENS,
  MODEL_COST_PER_1K_PROMPT_TOKENS,
} from "./anthropic.constants";
import { AnthropicMetrics } from "./anthropic.types";

@Injectable()
export class AnthropicProvider {
  constructor(private readonly opentelemetryProvider: OpentelemetryProvider) {}

  async listModels() {
    return Object.keys(MODEL_COST_PER_1K_PROMPT_TOKENS);
  }

  getConfiguredCosts() {
    return Object.keys(MODEL_COST_PER_1K_PROMPT_TOKENS).map((modelId) => {
      return {
        modelId,
        costPer1KPromptTokens: MODEL_COST_PER_1K_PROMPT_TOKENS[modelId],
        costPer1KCompletionTokens: MODEL_COST_PER_1K_COMPLETION_TOKENS[modelId],
      };
    });
  }

  getRequiredAttributesForPushMetrics() {
    return [
      {
        key: "promptTokens",
        valueType: "intValue",
      },
      {
        key: "completionTokens",
        valueType: "intValue",
      },
    ];
  }

  async saveMetrics(metrics: any, turoAttributes: any) {
    return this.opentelemetryProvider.saveMetrics(metrics, turoAttributes);
  }

  async fetchCostMetricsForSync(costEndpoint: any, queryOptions: any) {
    const result = await this.opentelemetryProvider.fetchCostMetricsForSync(
      costEndpoint,
      queryOptions,
    );
    const modelId = costEndpoint.config["modelId"];
    const transformedResult = {};

    Object.keys(result).forEach((_measurement) => {
      result[_measurement]["gauge"].forEach((dataPoint, index) => {
        if ("promptTokens" in dataPoint && "completionTokens" in dataPoint) {
          transformedResult[_measurement] = result[_measurement];
          const promptTokens = dataPoint["promptTokens"];
          const completionTokens = dataPoint["completionTokens"];
          const cost = this.getTokenCostForModel({
            modelId,
            promptTokens,
            completionTokens,
          }); // modelId will always be in supported model list as value is comming from dropdown in FE
          transformedResult[_measurement]["gauge"][index]["value"] = cost; //update numToken to cost
          transformedResult[_measurement]["gauge"][index]["modelId"] = modelId;
          transformedResult[_measurement]["gauge"][index]["type"] = "llm";
        }
      });
    });
    return transformedResult;
  }

  async encryptConfig(costConfig) {
    return costConfig;
  }

  async decryptConfig(costConfig) {
    return costConfig;
  }

  maskConfigSecret(costConfig) {
    return costConfig;
  }

  /**
   * Get the cost in USD for a given model and number of tokens.
   * @param metrics
   */
  getTokenCostForModel(metrics: AnthropicMetrics) {
    if (!(metrics.modelId in MODEL_COST_PER_1K_PROMPT_TOKENS)) {
      throw Error(
        `Unknown model: ${metrics.modelId}. Please provide a valid Anthropic model name.`,
      );
    }
    const promptTokenCost =
      MODEL_COST_PER_1K_PROMPT_TOKENS[metrics.modelId] *
      (metrics.promptTokens / 1000);
    const completionTokenCost =
      MODEL_COST_PER_1K_COMPLETION_TOKENS[metrics.modelId] *
      (metrics.completionTokens / 1000);
    return promptTokenCost + completionTokenCost;
  }
}