import Claude from "@anthropic-ai/sdk";
import { RESPONSE_LENGTH } from "./logic/constants";
import type { AIModelAPI, ChatMessage, InputToken } from "./types";
import { BOOKWORM_SYS } from "./prompts";
import type { AsyncRes } from "sortug";
import type {
  ImageBlockParam,
  MessageCreateParamsStreaming,
  TextBlockParam,
} from "@anthropic-ai/sdk/resources";

type Message = Claude.Messages.MessageParam;

export default class ClaudeAPI implements AIModelAPI {
  private model: string = "claude-opus-4-20250514";
  tokenizer: (text: string) => number;
  maxTokens: number;
  // model: string = "claude-3-5-sonnet-20241022";
  constructor(
    model: string,
    maxTokens = 200_000,
    tokenizer: (text: string) => number = (text) => text.length / 3,
  ) {
    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 buildInput(tokens: InputToken[]): Message[] {
    // can do base64 for images too
    const content = tokens.map((t) => {
      const content =
        "text" in t
          ? ({ type: "text", text: t.text } as TextBlockParam)
          : "img" in t
            ? ({
                type: "image",
                source: { type: "url", url: t.img },
              } as ImageBlockParam)
            : ({ type: "text", text: "oy vey" } as TextBlockParam);
      return content;
    });

    return [{ role: "user", content }];
  }

  // https://docs.anthropic.com/en/api/messages-examples#vision
  public async send(input: string | InputToken[], sys?: string) {
    const msgs: Message[] =
      typeof input === "string"
        ? [{ role: "user", content: input }]
        : this.buildInput(input);
    const truncated = this.truncateHistory(msgs);
    const res = await this.apiCall(truncated, sys);
    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([msg], sys);
    return res;
  }

  public async stream(
    input: string | InputToken[],
    handle: (c: any) => void,
    sys?: string,
  ) {
    const msgs: Message[] =
      typeof input === "string"
        ? [{ role: "user", content: input }]
        : this.buildInput(input);
    const truncated = this.truncateHistory(msgs);
    await this.apiCallStream(truncated, handle, sys);
  }

  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(
    messages: Message[],
    system?: string,
  ): Promise<AsyncRes<string>> {
    try {
      const claud = new Claude();
      const params = {
        model: this.model,
        max_tokens: RESPONSE_LENGTH,
        messages,
      };
      const res = await claud.messages.create(
        system ? { ...params, system } : params,
      );
      const resm: string = res.content.reduce((acc: string, item) => {
        if (item.type === "tool_use") return acc;
        else if (item.type === "text") return `${acc}\n${item.text}`;
        else return acc;
      }, "");
      // const resm = res.content.reduce((acc: string[], item) => {
      //     if (item.type === "tool_use") return acc;
      //     else if (item.type === "text") return [...acc, item.text]
      //     else return acc;
      //   }, []);
      return { ok: resm };
    } catch (e) {
      console.log(e, "error in claude api");
      return { error: `${e}` };
    }
  }

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

      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}`);
    }
  }
}