import Anthropic from "@anthropic-ai/sdk";
import { ChatCompletionResponse } from "@mistralai/mistralai";
import { LogDtoSourceEnum } from "@montelo/browser-client";
import OpenAI from "openai";
import { FC } from "react";
import { PrettyJson } from "~/components/PrettyJson";
import { CodeBlock } from "~/components/traces/MessagesView/CodeBlock";
import { RoleHeaderMap } from "~/components/traces/MessagesView/RoleHeaderMap";
import { Content, Message, MessagesContainer, ToolCallsComponent } from "~/components/traces/MessagesView/styles";
import { formatToolCalls } from "./utils";

type SourceToMessageMapType = {
  [key in LogDtoSourceEnum]: {
    InputComponent: FC<{ value: any; jsonView: boolean; sortAscending: boolean }>;
    OutputComponent: FC<{ value: any; jsonView: boolean; sortAscending: boolean }>;
  };
};

export const SourceToMessageMap: SourceToMessageMapType = {
  [LogDtoSourceEnum.Manual]: {
    InputComponent: ({ value, jsonView }) => {
      const RoleHeader = RoleHeaderMap.input;
      return (
        <Message>
          <RoleHeader />
          {jsonView ? (
            <CodeBlock value={JSON.stringify(value || {}, undefined, 2)} />
          ) : (
            <div className="bg-secondary rounded-xl px-3 py-1">
              <PrettyJson data={value} shouldTruncate={false} />
            </div>
          )}
        </Message>
      );
    },
    OutputComponent: ({ value, jsonView }) => {
      const RoleHeader = RoleHeaderMap.output;
      return (
        <Message>
          <RoleHeader />
          {jsonView ? (
            <CodeBlock value={JSON.stringify(value || {}, undefined, 2)} />
          ) : (
            <div className="bg-secondary rounded-xl px-3 py-1">
              <PrettyJson data={value} shouldTruncate={false} />
            </div>
          )}
        </Message>
      );
    },
  },
  [LogDtoSourceEnum.Openai]: {
    InputComponent: ({ value, jsonView, sortAscending }) => {
      const typedValue = value as OpenAI.ChatCompletionCreateParams;

      return (
        <MessagesContainer sortAscending={sortAscending}>
          {typedValue.messages.map((value, index) => {
            const RoleHeader = RoleHeaderMap[value.role];

            // @ts-expect-error
            const toolName = value.role === "tool" ? value.name : null;
            const toolCalls = "tool_calls" in value && value.tool_calls ? formatToolCalls(value.tool_calls) : [];

            const CodeView = <CodeBlock value={JSON.stringify(typedValue, undefined, 2)} />;
            const PrettyContent =
              typeof value.content === "string" && value.content.length ? (
                <Content>{value.content}</Content>
              ) : toolCalls?.length ? (
                <ToolCallsComponent toolCalls={toolCalls} />
              ) : (
                CodeView
              );

            return (
              <Message key={index}>
                <RoleHeader toolName={toolName} />
                {jsonView ? CodeView : PrettyContent}
              </Message>
            );
          })}
        </MessagesContainer>
      );
    },
    OutputComponent: ({ value, jsonView }) => {
      const typedValue = value as OpenAI.ChatCompletion;
      const message = typedValue.choices[0].message;
      const RoleHeader = RoleHeaderMap[message.role];
      const isToolCall = !message.content && !!message.tool_calls?.length;
      const isFunctionCall = !message.content && !!message.function_call;

      const toolCalls = isToolCall ? formatToolCalls(message.tool_calls) : [];

      const getContent = () => {
        if (isToolCall) {
          return <ToolCallsComponent toolCalls={toolCalls} />;
        } else if (isFunctionCall) {
          return <CodeBlock value={JSON.stringify(message.function_call, undefined, 2)}></CodeBlock>;
        }

        return <Content>{message.content}</Content>;
      };
      const ContentRender = getContent();

      return (
        <Message>
          <RoleHeader />
          {jsonView ? <CodeBlock value={JSON.stringify(value, undefined, 2)} /> : ContentRender}
        </Message>
      );
    },
  },
  [LogDtoSourceEnum.Mistral]: {
    InputComponent: ({ value, jsonView, sortAscending }) => (
      <MessagesContainer sortAscending={sortAscending}>
        {value.messages.map(({ role, content }: { role: string; content: string | string[] }, index: number) => {
          const RoleHeader = RoleHeaderMap[role];
          const strContent =
            typeof content === "string" ? content : content.reduce((acc, content) => acc + content, "");
          return (
            <Message key={index}>
              {<RoleHeader />}
              {jsonView ? <CodeBlock value={JSON.stringify(value, undefined, 2)} /> : <Content>{strContent}</Content>}
            </Message>
          );
        })}
      </MessagesContainer>
    ),
    OutputComponent: ({ value, jsonView }) => {
      const typedValue = value as ChatCompletionResponse;
      const message = typedValue.choices[0].message;
      const RoleHeader = RoleHeaderMap[message.role];
      const content = message.content;

      return (
        <Message>
          <RoleHeader />
          {jsonView ? <CodeBlock value={JSON.stringify(value, undefined, 2)} /> : <Content>{content}</Content>}
        </Message>
      );
    },
  },
  [LogDtoSourceEnum.Anthropic]: {
    InputComponent: ({ value, jsonView, sortAscending }) => {
      const typedValue = value as Anthropic.MessageCreateParams;

      return (
        <MessagesContainer sortAscending={sortAscending}>
          {typedValue.messages.map((value, index) => {
            const RoleHeader = RoleHeaderMap[value.role];

            const CodeView = <CodeBlock value={JSON.stringify(value, undefined, 2)} />;
            const PrettyContent = typeof value.content === "string" ? <Content>{value.content}</Content> : CodeView;

            return (
              <Message key={index}>
                <RoleHeader />
                {jsonView ? CodeView : PrettyContent}
              </Message>
            );
          })}
        </MessagesContainer>
      );
    },
    OutputComponent: ({ value, jsonView }) => {
      const typedValue = value as Anthropic.Message;
      const RoleHeader = RoleHeaderMap[typedValue.role];
      const stringOutput = typedValue.content.reduce((acc, content) => acc + content.text, "");

      return (
        <Message>
          <RoleHeader />
          {jsonView ? <CodeBlock value={JSON.stringify(value, undefined, 2)} /> : <Content>{stringOutput}</Content>}
        </Message>
      );
    },
  },
  [LogDtoSourceEnum.Cohere]: {
    InputComponent: () => <p>Unimplemented</p>,
    OutputComponent: () => <p>Unimplemented</p>,
  },
};
