import {
  Post,
  PostEmbed,
  RecordEmbed,
  RecordWithMediaEmbed,
} from "@skyware/bot";
import * as c from "../core";
import * as yaml from "js-yaml";
import type { ParsedPost } from "../types";
import { postCache } from "../utils/cache";
import { env } from "../env";

const MAX_ATTACHMENTS = 2;

export async function parsePost(
  post: Post,
  includeThread: boolean,
  seenUris: Set<string> = new Set(),
  media: string[] = [],
): Promise<ParsedPost | undefined> {
  if (seenUris.has(post.uri)) {
    return undefined;
  }
  seenUris.add(post.uri);

  // Attachments are numbered in the order posts are parsed. The numbers are
  // shared with the model so it can match each picture to the post it came from.
  const images = parsePostImages(post, media);
  const video = parsePostVideo(post, media);

  const [quotePost, thread] = await Promise.all([
    parseQuote(post, seenUris, media),
    includeThread ? parseThreadContext(post, seenUris, media) : undefined,
  ]);

  return {
    author: post.author.displayName
      ? `${post.author.displayName} (${post.author.handle})`
      : `Handle: ${post.author.handle}`,
    text: post.text,
    ...(post.langs?.length && { language: post.langs.join(", ") }),
    ...(images.length > 0 && { images }),
    ...(video && { video }),
    ...(quotePost && { quotePost }),
    ...(thread && { thread }),
  };
}

/**
 * Turns posts into the YAML context given to the model, along with the image
 * URLs that should be attached next to it.
 */
export async function buildContext(posts: Post[]) {
  const media: string[] = [];
  const seenUris = new Set<string>();
  const parsed: ParsedPost[] = [];

  // One at a time so attachment numbers follow the order of the posts
  for (const post of posts) {
    const result = await parsePost(post, true, seenUris, media);
    if (result) parsed.push(result);
  }

  return { media, context: yaml.dump({ posts: parsed }) };
}

/**
 * The parent chain of a reply, plus the root post when the chain is too long
 * to reach it, so the model always knows what the conversation started with.
 */
async function parseThreadContext(
  post: Post,
  seenUris: Set<string>,
  media: string[],
) {
  const chain = await traverseThread(post);

  const rootPost = await post.fetchRoot().catch(() => null);
  const rootMissing = rootPost !== null &&
    !chain.some((ancestor) => ancestor.uri === rootPost.uri);

  const ancestors = (await Promise.all(
    chain.map((ancestor) => parsePost(ancestor, false, seenUris, media)),
  )).filter((parsed): parsed is ParsedPost => parsed !== undefined);

  const root = rootMissing
    ? await parsePost(rootPost, false, seenUris, media)
    : undefined;

  if (ancestors.length === 0 && !root) return undefined;
  return {
    ...(root && { root, truncated: true }),
    ancestors,
  };
}

async function parseQuote(
  post: Post,
  seenUris: Set<string>,
  media: string[],
) {
  if (
    !post.embed || (!post.embed.isRecord() && !post.embed.isRecordWithMedia())
  ) return undefined;

  const record = (post.embed as RecordEmbed || RecordWithMediaEmbed).record;
  if (seenUris.has(record.uri)) {
    return undefined;
  }

  let embedPost = postCache.get(record.uri);
  if (!embedPost) {
    embedPost = await c.bot.getPost(record.uri);
    postCache.set(record.uri, embedPost);
  }

  return await parsePost(embedPost, false, seenUris, media);
}

function embeddedMedia(post: Post) {
  const embed = post.embed;
  if (!embed) return undefined;
  return embed.isRecordWithMedia() ? embed.media : embed;
}

/** Reserves an attachment slot, or returns undefined when the limit is reached. */
function attach(media: string[], url: string | undefined) {
  if (!url || !env.ENABLE_VISION || media.length >= MAX_ATTACHMENTS) {
    return undefined;
  }
  media.push(url);
  return media.length;
}

export function parsePostImages(post: Post, media: string[] = []) {
  const embed = embeddedMedia(post);
  if (!embed || !embed.isImages()) return [];

  return embed.images.map((image, idx) => {
    const attachment = attach(media, image.url);
    return {
      index: idx + 1,
      ...(image.alt && { alt: image.alt }),
      ...(attachment && { attachment }),
    };
  });
}

function parsePostVideo(post: Post, media: string[]) {
  const embed = embeddedMedia(post);
  if (!embed || !embed.isVideo()) return undefined;

  // Only the thumbnail can be shown to the model, not the video itself
  const attachment = attach(media, embed.thumbnail);
  return {
    ...(embed.alt && { alt: embed.alt }),
    ...(attachment && { thumbnailAttachment: attachment }),
  };
}

/*
  Traversal
*/
export async function traverseThread(post: Post): Promise<Post[]> {
  const thread: Post[] = [
    post,
  ];
  let currentPost: Post | undefined = post;
  let parentCount = 0;

  while (
    currentPost && parentCount < c.MAX_THREAD_DEPTH
  ) {
    const parentPost = await currentPost.fetchParent();

    if (parentPost) {
      thread.push(parentPost);
      currentPost = parentPost;
    } else {
      break;
    }
    parentCount++;
  }

  return thread.reverse();
}

export function parseThread(thread: Post[]) {
  return yaml.dump({
    uri: thread[0]!.uri,
    posts: thread.map((post) => ({
      author: `${post.author.displayName} (${post.author.handle})`,
      text: post.text,
    })),
  });
}