import Anthropic from "@anthropic-ai/sdk";
import { encoding_for_model, TiktokenModel } from "tiktoken";
import {
  toAnthropicMessages,
  getAPIKey,
  toOpenaiMessages,
  ApptileInMemoryChatMessage,
  sequenceToolCallAndResponses,
  sequenceToolCallsAndAddEmptyResults,
} from "../database/cliconfig";
import path from "path";
import fs from "fs";
import chalk from "chalk";
import { LLMProviders } from "./agents/baseAgent";
import { TokenUsage } from "@aws-sdk/client-bedrock-runtime";

// Function to count OpenAI tokens
function countOpenAITokens(messages: any[], model: TiktokenModel) {
  const encoder = encoding_for_model(model);

  let tokenCount = 0;
  for (const msg of messages) {
    tokenCount += encoder.encode(msg.role || "").length;
    if (msg.role === "assistant") {
      if (typeof msg.content === "string") {
        tokenCount += encoder.encode(msg.content).length;
      } else if (Array.isArray(msg.content)) {
        for (let block of msg.content) {
          if (block.type === "text") {
            tokenCount += encoder.encode(block.text).length;
          } else {
            throw new Error("Cannot encode non text tokens yet");
          }
        }
      }

      if (msg.tool_calls) {
        for (let toolCall of msg.tool_calls) {
          tokenCount += encoder.encode(toolCall.id).length;
          tokenCount += encoder.encode(toolCall.type).length;
          tokenCount += encoder.encode(toolCall.function.name).length;
          tokenCount += encoder.encode(toolCall.function.arguments).length;
        }
      }
    } else if (msg.role === "tool") {
      tokenCount += encoder.encode(msg.tool_call_id).length;
      if (typeof msg.content === "string") {
        tokenCount += encoder.encode(msg.content).length;
      } else {
        for (let block of msg.content) {
          if (block.type === "text") {
            tokenCount += encoder.encode(block.text).length;
          } else {
            throw new Error("Cannot encode non text tokens yet");
          }
        }
      }
    } else if (msg.role === "user") {
      if (typeof msg.content === "string") {
        tokenCount += encoder.encode(msg.content).length;
      } else {
        for (let block of msg.content) {
          if (block.type === "text") {
            tokenCount += encoder.encode(block.text).length;
          } else {
            throw new Error("Cannot encode non text tokens yet");
          }
        }
      }
    }
  }
  return tokenCount;
}

// Define a tool interface
export interface Tool {
  name: string;
  description: string;
  input_schema: any;
}

// Define the return type for token count and cost
export interface TokenCountResult {
  inputTokenCount: number;
  outputTokenCount: number;
  model: string;
  provider: LLMProviders;
  tools?: Tool[];
  agentType: string;
}

export interface TokenCalculationResult
  extends Omit<TokenCountResult, "inputTokenCount" | "outputTokenCount"> {
  tokenCount: number;
}

// Define the structure for tool usage information
export interface ToolUsage {
  tool_name: string;
  count: number;
  inputs: string[]; // Store the actual inputs used for each tool call
}

// Define the return type for tool usage analysis
export interface ToolUsageResult {
  totalToolCalls: number;
  uniqueTools: number;
  tools: ToolUsage[];
  chatId: number;
  model: string;
  provider: LLMProviders;
  type: "widget" | "app" | string;
}

export interface ITokenCounterRequest {
  model: string;
  provider: LLMProviders;
  inMemoryMessages: ApptileInMemoryChatMessage[];
  agentType: string;
}

// TokenCounter class
export class TokenCounter {
  private tools: Anthropic.Tool[];
  private inMemoryInputMessages: ApptileInMemoryChatMessage[][];
  private inMemoryOutputMessages: ApptileInMemoryChatMessage[];
  private agentType: string;
  private model: string;
  private provider: LLMProviders;
  private amazonUsage: TokenUsage[];

  /**
   * Creates a new TokenCounter instance
   * @param tools The tools to use for token counting
   */
  constructor(
    tools: Anthropic.Tool[] = [],
    agentType: string = "",
    model: string = "",
    provider: LLMProviders = "openai"
  ) {
    this.tools = tools;
    this.inMemoryInputMessages = [];
    this.agentType = agentType;
    this.model = model;
    this.provider = provider;
    this.inMemoryOutputMessages = [];
    this.amazonUsage = [];
  }

  /**
   * Sets the tools for this token counter
   * @param tools The tools to use
   */
  setTools(tools: Anthropic.Tool[]): void {
    this.tools = tools;
  }

  appendInMemoryInputMessages(messages: ApptileInMemoryChatMessage[]): void {
    if (this.provider !== "amazon") {
      this.inMemoryInputMessages.push(messages);
    }
  }

  appendInMemoryOutputMessages(messages: ApptileInMemoryChatMessage[]): void {
    if (this.provider !== "amazon") {
      this.inMemoryOutputMessages.push(...messages);
    }
  }

  async getInputTokenCount(): Promise<TokenCalculationResult> {
    if (this.provider === "amazon") {
      if (this.amazonUsage.length > 0) {
        let total = this.amazonUsage.reduce((totalInput, usage) => totalInput + (usage.cacheReadInputTokens || 0) + (usage.inputTokens || 0), 0);
        return {
          tokenCount: total,
          model: this.model,
          provider: this.provider,
          agentType: this.agentType,
        };
      } else {
        throw new Error("Uninitialized usage in amazon");
      }
    } else {
      // Get the system prompt based on the agent type
      let systemPrompt = "You are a helpful assistant.";

      if (!this.agentType) {
        throw new Error(`Agent type not specified`);
      }

      if (this.agentType === "widget") {
        systemPrompt = fs.readFileSync(
          path.resolve(__dirname, "../systemPrompts/latest.html"),
          { encoding: "utf8" }
        );
      } else if (this.agentType === "app") {
        systemPrompt = fs.readFileSync(
          path.resolve(__dirname, "../systemPrompts/appPlanner.html"),
          { encoding: "utf8" }
        );
      } else {
        throw new Error(`Unknown agent type: ${this.agentType}`);
      }

      let tokenCount: number;

      if (this.provider === "claude") {
        try {
          const rec = getAPIKey("claude");
          if (!rec) {
            throw new Error("Claude API key not found");
          }

          const claude = new Anthropic({ apiKey: rec.apikey });

          const tokenCounts: [number] = [0];

          await Promise.all(
            this.inMemoryInputMessages.map(async (messages) => {
              try {
                const tokenCountResult = await claude.messages.countTokens({
                  model: this.model,
                  system: systemPrompt,
                  messages: toAnthropicMessages(
                    sequenceToolCallsAndAddEmptyResults(messages)
                  ),
                  tools: this.tools,
                });
                tokenCounts.push(tokenCountResult.input_tokens);
              } catch (error) {
                console.error("Failed to count tokens for messages: ", error, JSON.stringify(messages, null, 2));
              }
            })
          );

          tokenCount = tokenCounts.reduce((acc, count) => acc + count, 0);
        } catch (err: any) {
          console.error("Failed to count input tokens", err);
          throw err;
        }
      } else {
        // For OpenAI, use our own implementation of token counting
        const tokenCounts: [number] = [0];
        await Promise.all(
          this.inMemoryInputMessages.map(async (messages) => {
            const openaiMsgs = toOpenaiMessages(messages, systemPrompt);
            const count = countOpenAITokens(
              openaiMsgs,
              this.model as TiktokenModel
            );
            tokenCounts.push(count);
          })
        );

        tokenCount = tokenCounts.reduce((acc, count) => acc + count, 0);

        // Add an estimate for tool definitions (very rough estimate)
        const toolsTokenEstimate = this.tools.reduce((acc, tool) => {
          // Rough estimate: name + description + JSON.stringify(inputSchema)
          return (
            acc +
            tool.name.length +
            (tool.description || "").length +
            JSON.stringify(tool.input_schema).length
          );
        }, 0);

        // Convert characters to tokens (rough estimate: 4 chars = 1 token)
        tokenCount += Math.ceil(toolsTokenEstimate / 4);
      }

      // Return token count, cost, and tools information
      return {
        tokenCount,
        model: this.model,
        provider: this.provider,
        agentType: this.agentType,
      };
    }
  }

  /**
   * Gets the token count for output messages
   * @returns The token count result for output messages
   */
  async getOutputTokenCount(): Promise<TokenCalculationResult> {
    if (this.provider === "amazon") {
      if (this.amazonUsage) {
        let total = this.amazonUsage.reduce((totalOutput, usage) => totalOutput + (usage.outputTokens || 0), 0);
        return {
          tokenCount: total,
          model: this.model,
          provider: this.provider,
          agentType: this.agentType,
        }
      } else {
        throw new Error("Uninitialized usage in amazon");
      }
    } else {
      // For output messages, we don't need system prompt or tools
      let tokenCount: number;

      if (this.provider === "claude") {
        try {
          const rec = getAPIKey("claude");
          if (!rec) {
            throw new Error("Claude API key not found");
          }

          const claude = new Anthropic({ apiKey: rec.apikey });

          const tokenCountResult = await claude.messages.countTokens({
            model: this.model,
            messages: toAnthropicMessages(
              sequenceToolCallAndResponses(this.inMemoryOutputMessages)
            ),
          });

          tokenCount = tokenCountResult.input_tokens;
        } catch (err: any) {
          console.error(
            "Failed to count output tokens",
            err,
            chalk.redBright(
              JSON.stringify(
                this.inMemoryOutputMessages.map(
                  (it: ApptileInMemoryChatMessage) => {
                    return {
                      role: it.role,
                      tool_call_id: it.tool_call_id,
                      content_type: it.content_type,
                    };
                  }
                ),
                null,
                2
              )
            )
          );
          throw err;
        }
      } else {
        // For OpenAI, use our own implementation of token counting
        const openaiMsgs = toOpenaiMessages(this.inMemoryOutputMessages, "");
        tokenCount = countOpenAITokens(openaiMsgs, this.model as TiktokenModel);
      }

      // Return token count, cost, and tools information
      return {
        tokenCount,
        model: this.model,
        provider: this.provider,
        agentType: this.agentType,
      };
    }
  }

  setAmazonUsage(usage: TokenUsage) {
    this.amazonUsage.push(usage);
  }

  /**
   * Gets the token count for both input and output messages
   * @returns The combined token count result
   */
  async getTokenCount(): Promise<TokenCountResult> {
    console.info(chalk.red("Token usage: " + JSON.stringify(this.amazonUsage, null, 2)));
    if (this.provider === "amazon" && this.amazonUsage) {
      const totalInput = this.amazonUsage.reduce((totalInput, usage) => totalInput + (usage.cacheReadInputTokens || 0) + (usage.inputTokens || 0), 0);
      const totalOutput = this.amazonUsage.reduce((totalOutput, usage) => totalOutput + (usage.outputTokens || 0), 0);
      
      return {
        inputTokenCount: totalInput,
        outputTokenCount: totalOutput,
        model: this.model,
        provider: this.provider,
        agentType: this.agentType
      };
    } else {
      try {
        const inputResult = await this.getInputTokenCount();
        const outputResult = await this.getOutputTokenCount();

        // Combine the results
        return {
          inputTokenCount: inputResult.tokenCount,
          outputTokenCount: outputResult.tokenCount,
          model: this.model,
          provider: this.provider,
          agentType: this.agentType,
        };
      } catch (err) {
        console.error("Failed to get token counts: ", err);
        return {
          inputTokenCount: 0,
          outputTokenCount: 0,
          model: "",
          provider: "claude",
          agentType: "",
        };
      }
    }
  }
}
