import { parsePartialJsonObject } from "assistant-stream/utils";
import { generateId } from "../../utils/id";
import type {
  ReasoningMessagePart,
  SourceMessagePart,
  ThreadStep,
  MessageStatus,
  ImageMessagePart,
  ThreadMessage,
  ThreadAssistantMessagePart,
  ThreadAssistantMessage,
  ThreadUserMessagePart,
  ThreadUserMessage,
  ThreadSystemMessage,
  FileMessagePart,
  DataMessagePart,
  Unstable_AudioMessagePart,
} from "../../types/message";
import type { CompleteAttachment } from "../../types/attachment";
import type { MessageTiming, TextMessagePart } from "../../types/message";
import type {
  ReadonlyJSONObject,
  ReadonlyJSONValue,
} from "assistant-stream/utils";

type DataPrefixedPart = {
  readonly type: `data-${string}`;
  readonly data: any;
};

export type ThreadMessageLike = {
  readonly role: "assistant" | "user" | "system";
  readonly content:
    | string
    | readonly (
        | TextMessagePart
        | ReasoningMessagePart
        | SourceMessagePart
        | ImageMessagePart
        | FileMessagePart
        | DataMessagePart
        | Unstable_AudioMessagePart
        | DataPrefixedPart
        | {
            readonly type: "tool-call";
            readonly toolCallId?: string;
            readonly toolName: string;
            readonly args?: ReadonlyJSONObject;
            readonly argsText?: string;
            readonly artifact?: any;
            readonly result?: any | undefined;
            readonly isError?: boolean | undefined;
            readonly parentId?: string | undefined;
            readonly messages?: readonly ThreadMessage[] | undefined;
          }
      )[];
  readonly id?: string | undefined;
  readonly createdAt?: Date | undefined;
  readonly status?: MessageStatus | undefined;
  readonly attachments?:
    | readonly (Omit<CompleteAttachment, "content"> & {
        readonly content: readonly (ThreadUserMessagePart | DataPrefixedPart)[];
      })[]
    | undefined;
  readonly metadata?:
    | {
        readonly unstable_state?: ReadonlyJSONValue;
        readonly unstable_annotations?:
          | readonly ReadonlyJSONValue[]
          | undefined;
        readonly unstable_data?: readonly ReadonlyJSONValue[] | undefined;
        readonly steps?: readonly ThreadStep[] | undefined;
        readonly timing?: MessageTiming | undefined;
        readonly submittedFeedback?: { readonly type: "positive" | "negative" };
        readonly custom?: Record<string, unknown> | undefined;
      }
    | undefined;
};

const convertDataPrefixedPart = (
  type: string,
  data: unknown,
): DataMessagePart | undefined => {
  if (!type.startsWith("data-")) return undefined;
  return { type: "data", name: type.substring(5), data };
};

export const fromThreadMessageLike = (
  like: ThreadMessageLike,
  fallbackId: string,
  fallbackStatus: MessageStatus,
): ThreadMessage => {
  const { role, id, createdAt, attachments, status, metadata } = like;
  const common = {
    id: id ?? fallbackId,
    createdAt: createdAt ?? new Date(),
  };

  const content =
    typeof like.content === "string"
      ? [{ type: "text" as const, text: like.content }]
      : like.content;

  const sanitizeImageContent = ({
    image,
    ...rest
  }: ImageMessagePart): ImageMessagePart | null => {
    const match = image.match(
      /^data:image\/(png|jpeg|jpg|gif|webp);base64,(.*)$/,
    );
    if (match) {
      return { ...rest, image };
    }
    console.warn(`Invalid image data format detected`);
    return null;
  };

  if (role !== "user" && attachments?.length)
    throw new Error("attachments are only supported for user messages");

  if (role !== "assistant" && status)
    throw new Error("status is only supported for assistant messages");

  if (role !== "assistant" && metadata?.steps)
    throw new Error("metadata.steps is only supported for assistant messages");

  switch (role) {
    case "assistant":
      return {
        ...common,
        role,
        content: content
          .map((part): ThreadAssistantMessagePart | null => {
            const type = part.type;
            switch (type) {
              case "text":
              case "reasoning":
                if (part.text.trim().length === 0) return null;
                return part;

              case "file":
              case "source":
                return part;

              case "image":
                return sanitizeImageContent(part);

              case "data":
                return part;

              case "tool-call": {
                const { parentId, messages, ...basePart } = part;
                const commonProps = {
                  ...basePart,
                  toolCallId: part.toolCallId ?? `tool-${generateId()}`,
                  ...(parentId !== undefined && { parentId }),
                  ...(messages !== undefined && { messages }),
                };

                if (part.args) {
                  return {
                    ...commonProps,
                    args: part.args,
                    argsText: part.argsText ?? JSON.stringify(part.args),
                  };
                }
                return {
                  ...commonProps,
                  args: parsePartialJsonObject(part.argsText ?? "") ?? {},
                  argsText: part.argsText ?? "",
                };
              }

              default: {
                const converted = convertDataPrefixedPart(
                  type,
                  (part as DataPrefixedPart).data,
                );
                if (converted) return converted;
                throw new Error(
                  `Unsupported assistant message part type: ${type}`,
                );
              }
            }
          })
          .filter((c) => !!c),
        status: status ?? fallbackStatus,
        metadata: {
          unstable_state: metadata?.unstable_state ?? null,
          unstable_annotations: metadata?.unstable_annotations ?? [],
          unstable_data: metadata?.unstable_data ?? [],
          custom: metadata?.custom ?? {},
          steps: metadata?.steps ?? [],
          ...(metadata?.timing && { timing: metadata.timing }),
          ...(metadata?.submittedFeedback && {
            submittedFeedback: metadata.submittedFeedback,
          }),
        },
      } satisfies ThreadAssistantMessage;

    case "user":
      return {
        ...common,
        role,
        content: content.map((part): ThreadUserMessagePart => {
          const type = part.type;
          switch (type) {
            case "text":
            case "image":
            case "audio":
            case "file":
            case "data":
              return part;

            default: {
              const converted = convertDataPrefixedPart(
                type,
                (part as DataPrefixedPart).data,
              );
              if (converted) return converted;
              throw new Error(`Unsupported user message part type: ${type}`);
            }
          }
        }),
        attachments: (attachments ?? []).map((att) => ({
          ...att,
          content: att.content.map((part): ThreadUserMessagePart => {
            const converted = convertDataPrefixedPart(
              part.type,
              (part as DataPrefixedPart).data,
            );
            return converted ?? (part as ThreadUserMessagePart);
          }),
        })),
        metadata: {
          custom: metadata?.custom ?? {},
        },
      } satisfies ThreadUserMessage;

    case "system":
      if (content.length !== 1 || content[0]!.type !== "text")
        throw new Error(
          "System messages must have exactly one text message part.",
        );

      return {
        ...common,
        role,
        content: content as [TextMessagePart],
        metadata: {
          custom: metadata?.custom ?? {},
        },
      } satisfies ThreadSystemMessage;

    default: {
      const unsupportedRole: never = role;
      throw new Error(`Unknown message role: ${unsupportedRole}`);
    }
  }
};