import type { Presentation } from "./utils";

// AES-256-GCM helpers using Web Crypto API
// Each upload gets a unique key, so fixed IVs are safe (one use per key).
// Metadata (filename + size) is encrypted with IV 0.
// File body is split into 16K chunks; IV starts at 2, increments by 2.
// LSB of final chunk's IV is set to 1 to authenticate stream completeness.

const META_IV = new Uint8Array(12); // iv 0 — 12 zero bytes

export const CHUNK_SIZE = 16384; // 16K plaintext per chunk
export const CHUNK_CT_SIZE = CHUNK_SIZE + 16; // 16K + GCM tag

// IV is a 12-byte big-endian integer; counter stored in the last 4 bytes.
export function chunkIv(counter: number): Uint8Array<ArrayBuffer> {
  const iv = new Uint8Array(12);
  new DataView(iv.buffer).setUint32(8, counter);
  return iv;
}

export async function generateKey() {
  const key = await crypto.subtle.generateKey(
    { name: "AES-GCM", length: 256 },
    true,
    ["encrypt", "decrypt"],
  );
  const raw = await crypto.subtle.exportKey("raw", key);
  return {
    key,
    encoded: new Uint8Array(raw).toBase64({
      alphabet: "base64url",
      omitPadding: true,
    }),
  };
}

export async function importKey(encoded: string) {
  const raw = Uint8Array.fromBase64(encoded, { alphabet: "base64url" });
  return crypto.subtle.importKey("raw", raw, { name: "AES-GCM" }, false, [
    "encrypt",
    "decrypt",
  ]);
}

// Metadata plaintext is padded with trailing spaces to this length,
// so the ciphertext (plaintext + 16-byte GCM tag) is always exactly 4096 bytes.
const META_PLAIN_LEN = 4096 - 16;

// Encrypt metadata JSON { name, size, presentation } with IV 0, padded to META_PLAIN_LEN bytes
export async function encryptMeta(
  fileName: string,
  fileLen: number,
  presentation: Presentation,
  key: CryptoKey,
) {
  const nameBytes = new TextEncoder().encode(fileName);
  if (nameBytes.length > 2048) throw new Error("Filename too long");
  const json = JSON.stringify({ name: fileName, size: fileLen, presentation });
  const jsonBytes = new TextEncoder().encode(json);
  if (jsonBytes.length > META_PLAIN_LEN) throw new Error("Metadata too large");
  // Pad to exactly META_PLAIN_LEN with trailing spaces (valid JSON whitespace)
  const plain = new Uint8Array(META_PLAIN_LEN);
  plain.set(jsonBytes);
  plain.fill(0x20, jsonBytes.length); // 0x20 = space
  const ct = await crypto.subtle.encrypt(
    { name: "AES-GCM", iv: META_IV },
    key,
    plain,
  );
  return new Uint8Array(ct);
}

// Decrypt metadata with IV 0 — returns { fileName, fileLen }
export async function decryptMeta(
  ciphertext: Uint8Array<ArrayBuffer>,
  key: CryptoKey,
) {
  const plain = await crypto.subtle.decrypt(
    { name: "AES-GCM", iv: META_IV },
    key,
    ciphertext,
  );
  const { name, size, presentation } = JSON.parse(new TextDecoder().decode(plain));
  return { fileName: name as string, fileLen: size as number, presentation: presentation as Presentation };
}

// Encrypt file body as 16K chunks; IVs 2, 4, 6, … with LSB set on final chunk
export async function encrypt(fileBuffer: ArrayBuffer, key: CryptoKey) {
  const numChunks = Math.max(Math.ceil(fileBuffer.byteLength / CHUNK_SIZE), 1);
  const padded = new Uint8Array(numChunks * CHUNK_SIZE);
  padded.set(new Uint8Array(fileBuffer));

  const result = new Uint8Array(numChunks * CHUNK_CT_SIZE);
  for (let i = 0; i < numChunks; i++) {
    const isFinal = i === numChunks - 1;
    const iv = chunkIv((i + 1) * 2 | (isFinal ? 1 : 0));
    const chunkBuf = padded.buffer.slice(i * CHUNK_SIZE, (i + 1) * CHUNK_SIZE);
    const ct = await crypto.subtle.encrypt({ name: "AES-GCM", iv }, key, chunkBuf);
    result.set(new Uint8Array(ct), i * CHUNK_CT_SIZE);
  }
  return result;
}

function numChunks(fileLen: number) {
  return Math.max(Math.ceil(fileLen / CHUNK_SIZE), 1);
}

async function decryptChunk(
  ct: ArrayBuffer,
  chunkIndex: number,
  fileLen: number,
  key: CryptoKey,
): Promise<Uint8Array> {
  const isFinal = chunkIndex === numChunks(fileLen) - 1;
  const iv = chunkIv((chunkIndex + 1) * 2 | (isFinal ? 1 : 0));
  const plain = await crypto.subtle.decrypt({ name: "AES-GCM", iv }, key, ct);
  const bytes = new Uint8Array(plain);
  return isFinal ? bytes.slice(0, fileLen - chunkIndex * CHUNK_SIZE) : bytes;
}

// Map a plaintext byte range to the ciphertext byte range needed to decrypt it.
export function ctRange(plainStart: number, plainEnd: number): [number, number] {
  const startChunk = Math.floor(plainStart / CHUNK_SIZE);
  const endChunk = Math.floor(plainEnd / CHUNK_SIZE);
  return [startChunk * CHUNK_CT_SIZE, (endChunk + 1) * CHUNK_CT_SIZE - 1];
}

// Streaming decryption of a plaintext byte range from a ciphertext stream.
// The reader must provide ciphertext starting at the chunk aligned to plainStart.
export function decryptRange(
  reader: ReadableStreamDefaultReader<Uint8Array>,
  key: CryptoKey,
  fileLen: number,
  plainStart: number,
  plainEnd: number,
): ReadableStream<Uint8Array> {
  const startChunk = Math.floor(plainStart / CHUNK_SIZE);
  const endChunk = Math.floor(plainEnd / CHUNK_SIZE);

  return new ReadableStream({
    async start(controller) {
      let buf = new Uint8Array(0);
      let chunkIndex = startChunk;

      try {
        while (true) {
          const { done, value } = await reader.read();

          if (value) {
            const next = new Uint8Array(buf.length + value.length);
            next.set(buf);
            next.set(value, buf.length);
            buf = next;
          }

          while (buf.length >= CHUNK_CT_SIZE && chunkIndex <= endChunk) {
            const chunkBuf = buf.buffer.slice(
              buf.byteOffset,
              buf.byteOffset + CHUNK_CT_SIZE,
            );
            buf = new Uint8Array(buf.buffer.slice(buf.byteOffset + CHUNK_CT_SIZE));

            let plainBytes = await decryptChunk(chunkBuf, chunkIndex, fileLen, key);

            // Trim first/last chunks to the requested range
            const chunkPlainStart = chunkIndex * CHUNK_SIZE;
            const sliceStart = Math.max(0, plainStart - chunkPlainStart);
            const sliceEnd = Math.min(plainBytes.length, plainEnd - chunkPlainStart + 1);
            if (sliceStart > 0 || sliceEnd < plainBytes.length) {
              plainBytes = plainBytes.subarray(sliceStart, sliceEnd);
            }

            controller.enqueue(plainBytes);
            chunkIndex++;
          }

          if (done) break;
        }
        controller.close();
      } catch (err) {
        try { controller.error(err); } catch { /* already closed/cancelled */ }
      }
    },
    cancel() {
      reader.cancel();
    },
  });
}