import Claude from "@anthropic-ai/sdk";
import { RESPONSE_LENGTH } from "./logic/constants";
import type {
  AIModelAPI,
  ChatMessage,
  OChoice,
  OChunk,
  OMessage,
} from "./types";
import { BOOKWORM_SYS } from "./prompts";
import type { AsyncRes } from "sortug";

type Message = Claude.Messages.MessageParam;

export default class ClaudeAPI implements AIModelAPI {
  private model: string = "claude-3-7-sonnet-20250219";
  tokenizer: (text: string) => number;
  maxTokens: number;
  // model: string = "claude-3-5-sonnet-20241022";
  constructor(
    maxTokens = 200_000,
    tokenizer: (text: string) => number = (text) => text.length / 3,
    model?: string,
  ) {
    this.maxTokens = maxTokens;
    this.tokenizer = tokenizer;
    if (model) this.model = model;
  }
  public setModel(model: string) {
    this.model = model;
  }
  private mapMessages(input: ChatMessage[]): Message[] {
    return input.map((m) => {
      const role = m.author === "claude" ? "assistant" : "user";
      return { role, content: m.text };
    });
  }

  private mapMessagesR1(input: ChatMessage[]): Message[] {
    return input.reduce((acc: Message[], m, i) => {
      const prev = acc[i - 1];
      const role: any = m.author === "claude" ? "assistant" : "user";
      const msg = { role, content: m.text };
      if (prev?.role === role) acc[i - 1] = msg;
      else acc = [...acc, msg];
      return acc;
    }, []);
  }

  public async send(sys: string, input: ChatMessage[]) {
    const messages = this.mapMessages(input);
    const truncated = this.truncateHistory(messages);
    const res = await this.apiCall(sys, truncated);
    return res;
  }

  public async sendR1(input: ChatMessage[]) {
    const messages = this.mapMessagesR1(input);
    const truncated = this.truncateHistory(messages);
    const res = await this.apiCall("", truncated, true);
    return res;
  }
  public async sendDoc(data: string) {
    const sys = BOOKWORM_SYS;
    const msg: Message = {
      role: "user",
      content: [
        {
          type: "document",
          source: { type: "base64", data, media_type: "application/pdf" },
        },
        {
          type: "text",
          text: "Please analyze this according to your system prompt. Be thorough.",
        },
      ],
    };
    const res = await this.apiCall(sys, [msg]);
    return res;
  }

  public async stream(
    sys: string,
    input: ChatMessage[],
    handle: (c: any) => void,
  ) {
    const messages = this.mapMessages(input);
    const truncated = this.truncateHistory(messages);
    await this.apiCallStream(sys, truncated, handle);
  }

  public async streamR1(input: ChatMessage[], handle: (c: any) => void) {
    const messages = this.mapMessagesR1(input);
    const truncated = this.truncateHistory(messages);
    await this.apiCallStream("", truncated, handle, true);
  }

  private truncateHistory(messages: Message[]): Message[] {
    const totalTokens = messages.reduce((total, message) => {
      return total + this.tokenizer(message.content as string);
    }, 0);
    while (totalTokens > this.maxTokens && messages.length > 1) {
      messages.splice(0, 1);
    }
    return messages;
  }

  // TODO
  // https://docs.anthropic.com/en/api/messages-examples#putting-words-in-claudes-mouth
  private async apiCall(
    system: string,
    messages: Message[],
    isR1: boolean = false,
  ): Promise<AsyncRes<string[]>> {
    try {
      const claud = new Claude();
      // const list = await claud.models.list();
      // console.log(list.data);
      const res = await claud.messages.create({
        model: this.model,
        max_tokens: RESPONSE_LENGTH,
        system,
        messages,
      });
      return {
        ok: res.content.reduce((acc: string[], item) => {
          if (item.type === "tool_use") return acc;
          else return [...acc, item.text];
        }, []),
      };
    } catch (e) {
      console.log(e, "error in claude api");
      return { error: `${e}` };
    }
  }

  private async apiCallStream(
    system: string,
    messages: Message[],
    handle: (c: any) => void,
    isR1: boolean = false,
  ): Promise<void> {
    try {
      const claud = new Claude();
      const stream = await claud.messages.create({
        model: this.model,
        max_tokens: RESPONSE_LENGTH,
        system,
        messages,
        stream: true,
      });

      for await (const part of stream) {
        if (part.type === "message_start") continue;
        if (part.type === "content_block_start") continue;
        if (part.type === "content_block_delta") {
          console.log("delta", part.delta);
          const delta: any = part.delta;
          handle(delta.text);
        }
      }
    } catch (e) {
      console.log(e, "error in claude api");
      handle(`Error streaming Claude, ${e}`);
    }
  }
}