import path from "node:path";
import type {
  AdapterContext,
  ArtifactStore,
  ComputerRef,
  SandboxProvider,
} from "@rakazo/adapter-kit";
import { ArtifactTooLargeError } from "@rakazo/adapter-kit";
import type { ComputerMode, MessageBlock } from "@rakazo/contracts";
import { ATTACHMENT_MAX_BYTES } from "@rakazo/contracts";
import {
  attachmentExtensionForMimeType,
  inferAttachmentMimeType,
  messageBlockForArtifact,
} from "@rakazo/core";
import type { PrismaClient } from "@rakazo/db";
import { artifactMaxBytes, storeArtifactStream } from "./artifact-files.js";
import { discardByteStream } from "./byte-streams.js";
import {
  assertTeamWorkspacePathAllowed,
  computerFileTransferMaxBytes,
  resolveBotWorkspacePath,
  writeComputerFileStream,
} from "./computer-support.js";

export type MaterializedThreadFile = {
  artifactId: string;
  name: string;
  mimeType: string;
  size: number;
  /** Workspace path, or null when the file is too large to copy to this computer. */
  path: string | null;
};

type AttachmentBlock = Extract<MessageBlock, { kind: "image" | "video" | "audio" | "file" }>;

const ATTACHMENT_BLOCK_KINDS: ReadonlySet<string> = new Set(["image", "video", "audio", "file"]);

/** Media up to its kind cap when the store streams; otherwise the inline attachment cap. */
export function attachableFileMaxBytes(artifacts: ArtifactStore, mimeType: string): number {
  return artifacts.describe().capabilities.stream
    ? artifactMaxBytes(mimeType)
    : ATTACHMENT_MAX_BYTES;
}

export function formatMiB(bytes: number): string {
  return `${Math.floor(bytes / (1024 * 1024))} MiB`;
}

export async function attachWorkspaceFileToThread(
  deps: {
    prisma: PrismaClient;
    artifacts: ArtifactStore;
  },
  input: {
    spaceId: string;
    userId: string;
    botId: string;
    groupId?: string;
    runId: string;
    filePath: string;
    /** A stream (e.g. from readComputerFileStream) is stored without buffering the file. */
    content: Uint8Array | ReadableStream<Uint8Array>;
    operationId: string;
    signal?: AbortSignal;
    present?: boolean;
  },
): Promise<{ artifactId: string; block: AttachmentBlock }> {
  const content = input.content;
  const stream =
    content instanceof Uint8Array
      ? (async function* () {
          yield content;
        })()
      : content;
  const name = path.basename(input.filePath) || input.filePath;
  let row: Awaited<ReturnType<typeof storeArtifactStream>>;
  try {
    const mimeType = inferAttachmentMimeType(name);
    if (!mimeType) {
      throw new Error(`Unsupported attachment type for ${name}`);
    }
    row = await storeArtifactStream(deps, {
      spaceId: input.spaceId,
      userId: input.userId,
      botId: input.botId,
      groupId: input.groupId ?? null,
      runId: input.runId,
      name,
      mimeType,
      stream,
      maxBytes: attachableFileMaxBytes(deps.artifacts, mimeType),
      operationId: input.operationId,
      signal: input.signal,
    });
  } catch (error) {
    discardByteStream(stream);
    if (error instanceof ArtifactTooLargeError) {
      throw new Error(`file exceeds the ${formatMiB(error.maxBytes)} attachment limit`);
    }
    throw error;
  }
  const block = messageBlockForArtifact(row);
  const presentable = block.kind === "image" || block.kind === "file";
  return {
    artifactId: row.id,
    block: input.present && presentable ? { ...block, present: true } : block,
  };
}

/**
 * Copies this turn's attachments (files, images, video and audio) into the bot's workspace so its
 * tools can open them. Files above the computer's transfer cap stay in storage and are listed
 * with their artifact id instead.
 */
export async function materializeCurrentTurnFiles(
  deps: {
    prisma: PrismaClient;
    artifacts: ArtifactStore;
    sandbox: SandboxProvider;
  },
  blocks: MessageBlock[] | undefined,
  input: {
    context: AdapterContext & { botId: string };
    computer: ComputerRef;
    computerMode: ComputerMode;
    teamBotIds?: ReadonlySet<string>;
    markWorkspaceDirty?: () => void;
  },
): Promise<MaterializedThreadFile[]> {
  const attachmentBlocks = blocks?.filter((block): block is AttachmentBlock =>
    ATTACHMENT_BLOCK_KINDS.has(block.kind),
  );
  if (!attachmentBlocks?.length) return [];

  const rows = await deps.prisma.artifact.findMany({
    where: {
      id: { in: attachmentBlocks.map((block) => block.artifactId) },
      spaceId: input.context.spaceId,
      userId: input.context.userId,
    },
  });
  const byId = new Map(rows.map((row) => [row.id, row]));
  const transferMaxBytes = computerFileTransferMaxBytes(input.computer);
  const materialized: MaterializedThreadFile[] = [];

  for (const block of attachmentBlocks) {
    const row = byId.get(block.artifactId);
    if (!row) throw new Error(`Attached file is unavailable: ${block.name}`);
    const entry = { artifactId: row.id, name: row.name, mimeType: row.mimeType, size: row.size };
    if (row.size > Math.min(transferMaxBytes, artifactMaxBytes(row.mimeType))) {
      materialized.push({ ...entry, path: null });
      continue;
    }
    const relativePath = `attachments/${row.id}${attachmentExtensionForMimeType(row.mimeType)}`;
    const storedPath = resolveBotWorkspacePath(
      input.computerMode,
      input.context.botId,
      relativePath,
    );
    if (input.teamBotIds) assertTeamWorkspacePathAllowed(storedPath, input.teamBotIds);
    // Streamed, so a large video copied to a local computer never sits in memory whole.
    const stream = await deps.artifacts.createReadStream(row.storageKey, input.context);
    input.markWorkspaceDirty?.();
    await writeComputerFileStream(
      deps.sandbox,
      input.computer,
      { path: storedPath, stream },
      input.context,
    );
    materialized.push({ ...entry, path: relativePath });
  }

  return materialized;
}

export function currentTurnFilesInstruction(files: readonly MaterializedThreadFile[]): string {
  const copied = files.filter((file) => file.path !== null);
  const kept = files.filter((file) => file.path === null);
  const describe = (file: MaterializedThreadFile) =>
    `- ${JSON.stringify(file.name)} (${file.mimeType}, ${file.size} bytes)`;
  return [
    ...(copied.length
      ? [
          "Current-turn attached files are available on your computer. Inspect them before answering:",
          ...copied.map((file) => `${describe(file)}: ${JSON.stringify(file.path)}`),
        ]
      : []),
    ...(kept.length
      ? [
          "These current-turn attachments are too large to copy to your computer. Refer to them by artifact id; do not guess their contents:",
          ...kept.map((file) => `${describe(file)}: artifact ${file.artifactId}`),
        ]
      : []),
  ].join("\n");
}