// Blob handler tests - upload limits, scope, mime sniffing, download headers
import { describe, expect, it } from 'vitest';
import { createBlobHandlers } from '../packages/core/src/handlers/xrpc-blob.js';
import { cidToString, createBlobCid } from '../packages/core/src/repo.js';
import { createSha256 } from '../packages/core/src/sha256.js';

const DID = 'did:plc:blobtest';
const PNG = new Uint8Array([
  0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2,
]);

/**
 * @param {{did?: string|null, readOnly?: boolean, blobUploadLimit?: number, streaming?: boolean, hasher?: import('../packages/core/src/ports.js').HasherFactory}} [options]
 */
function createBlobs({
  did = DID,
  readOnly = false,
  blobUploadLimit = 5_000_000,
  streaming = false,
  hasher,
} = {}) {
  /** @type {Map<string, {mimeType: string, size: number}>} */
  const meta = new Map();
  /** @type {Map<string, {data: Uint8Array, mimeType: string}>} */
  const store = new Map();

  /** @type {Map<string, string[]>} blob cid -> record URIs referencing it */
  const links = new Map();

  const actorStorage = {
    meta,
    links,
    async putBlob(
      /** @type {string} */ cid,
      /** @type {string} */ mimeType,
      /** @type {number} */ size,
    ) {
      meta.set(cid, { mimeType, size });
    },
    async listBlobs(
      /** @type {string|null} */ cursor,
      /** @type {number} */ limit,
    ) {
      const all = [...meta.keys()].sort();
      const from = cursor ? all.indexOf(cursor) + 1 : 0;
      const page = all.slice(from, from + limit);
      const next = from + limit < all.length ? page[page.length - 1] : null;
      return { cids: page, cursor: next };
    },
    async listBlobRecordUris(/** @type {string} */ cid) {
      return links.get(cid) ?? [];
    },
    async getBlob(/** @type {string} */ cid) {
      const held = meta.get(cid);
      return held ? { ...held, createdAt: 0 } : null;
    },
  };

  /** @type {{writes: number[], committed: {cid: string, mimeType: string}|null, aborted: boolean}} */
  const writerLog = { writes: [], committed: null, aborted: false };

  /** @type {Array<{start: number, end: number}|null>} */
  const streamLog = [];

  const blobs = {
    store,
    async put(
      /** @type {string} */ _did,
      /** @type {string} */ cid,
      /** @type {Uint8Array} */ data,
      /** @type {string} */ mimeType,
    ) {
      store.set(cid, { data, mimeType });
    },
    async get(/** @type {string} */ _did, /** @type {string} */ cid) {
      return store.get(cid) || null;
    },
    ...(streaming
      ? {
          async getStream(
            /** @type {string} */ _did,
            /** @type {string} */ cid,
            /** @type {{start: number, end: number}|undefined} */ range,
          ) {
            streamLog.push(range ?? null);
            const held = store.get(cid);
            if (!held) return null;
            const data = range
              ? held.data.subarray(range.start, range.end + 1)
              : held.data;
            return {
              body: /** @type {ReadableStream} */ (
                new Response(/** @type {any} */ (data)).body
              ),
              mimeType: held.mimeType,
            };
          },
          async createWriter(/** @type {string} */ _did) {
            /** @type {Uint8Array[]} */
            const staged = [];
            return {
              async write(/** @type {Uint8Array} */ chunk) {
                writerLog.writes.push(chunk.length);
                staged.push(chunk);
              },
              async commit(
                /** @type {string} */ cid,
                /** @type {string} */ mimeType,
              ) {
                writerLog.committed = { cid, mimeType };
                const size = staged.reduce((sum, c) => sum + c.length, 0);
                const data = new Uint8Array(size);
                let at = 0;
                for (const c of staged) {
                  data.set(c, at);
                  at += c.length;
                }
                store.set(cid, { data, mimeType });
              },
              async abort() {
                writerLog.aborted = true;
              },
            };
          },
        }
      : {}),
  };

  const handlers = createBlobHandlers({
    actorStorage: /** @type {any} */ (actorStorage),
    blobs: /** @type {any} */ (blobs),
    blobUploadLimit,
    hasher,
    readOnly,
    getDid: async () => did,
    readOnlyError: () =>
      Response.json(
        { error: 'AuthenticationRequired', message: 'This PDS is read-only' },
        { status: 401 },
      ),
  });
  return { handlers, actorStorage, blobs, meta, store, writerLog, streamLog };
}

/**
 * @param {ReturnType<typeof createBlobHandlers>} handlers
 * @param {string} path
 * @param {{auth?: {did: string, scope: string}|null, body?: Uint8Array, contentType?: string, headers?: Record<string, string>}} [options]
 */
async function call(
  handlers,
  path,
  { auth = null, body, contentType, headers } = {},
) {
  const route = handlers.routes[path.split('?')[0]];
  const url = new URL(`https://pds.example.com${path}`);
  /** @type {RequestInit} */
  const init = { method: body ? 'POST' : 'GET' };
  if (body) init.body = /** @type {BodyInit} */ (body);
  if (contentType || headers) {
    init.headers = {
      ...(headers ?? {}),
      ...(contentType ? { 'content-type': contentType } : {}),
    };
  }
  const request = new Request(url, init);
  const handler =
    /** @type {(request: Request, url: URL, auth: unknown) => Promise<Response>} */ (
      route.handler
    );
  return handler(request, url, auth);
}

const UPLOAD = '/xrpc/com.atproto.repo.uploadBlob';
const GET = '/xrpc/com.atproto.sync.getBlob';
const LIST = '/xrpc/com.atproto.sync.listBlobs';
const FULL = { did: DID, scope: 'com.atproto.access' };

describe('uploadBlob', () => {
  it('stores the blob and answers with its ref', async () => {
    const { handlers, meta, store } = createBlobs();
    const response = await call(handlers, UPLOAD, { auth: FULL, body: PNG });
    const body = await response.json();

    expect(response.status).toBe(200);
    expect(body.blob.$type).toBe('blob');
    expect(body.blob.mimeType).toBe('image/png');
    expect(body.blob.size).toBe(PNG.length);

    const cid = body.blob.ref.$link;
    expect(meta.get(cid)).toEqual({ mimeType: 'image/png', size: PNG.length });
    expect(store.get(cid)?.data).toEqual(PNG);
  });

  it('reads the type from the content, not the header', async () => {
    const { handlers } = createBlobs();
    const response = await call(handlers, UPLOAD, {
      auth: FULL,
      body: PNG,
      contentType: 'text/plain',
    });

    expect((await response.json()).blob.mimeType).toBe('image/png');
  });

  it('falls back to the header when the content says nothing', async () => {
    const { handlers } = createBlobs();
    const response = await call(handlers, UPLOAD, {
      auth: FULL,
      body: new Uint8Array([1, 2, 3, 4]),
      contentType: 'application/x-thing',
    });

    expect((await response.json()).blob.mimeType).toBe('application/x-thing');
  });

  it('needs a session', async () => {
    const { handlers } = createBlobs();
    const response = await call(handlers, UPLOAD, { body: PNG });

    expect(response.status).toBe(401);
    expect((await response.json()).error).toBe('AuthenticationRequired');
  });

  it('refuses a read-only server', async () => {
    const { handlers } = createBlobs({ readOnly: true });
    const response = await call(handlers, UPLOAD, { auth: FULL, body: PNG });

    expect(response.status).toBe(401);
    expect((await response.json()).message).toMatch(/read-only/);
  });

  it('needs an account', async () => {
    const { handlers } = createBlobs({ did: null });
    const response = await call(handlers, UPLOAD, { auth: FULL, body: PNG });

    expect(response.status).toBe(400);
    expect((await response.json()).message).toMatch(/not initialized/);
  });

  it('reports the limit when the blob is too large', async () => {
    const { handlers } = createBlobs({ blobUploadLimit: 4 });
    const response = await call(handlers, UPLOAD, { auth: FULL, body: PNG });
    const body = await response.json();

    expect(response.status).toBe(413);
    expect(body.error).toBe('PayloadTooLarge');
    expect(body.message).toContain('maximum size of 4 bytes');
  });

  it('refuses a type the scope does not allow', async () => {
    const { handlers, store } = createBlobs();
    const response = await call(handlers, UPLOAD, {
      auth: { did: DID, scope: 'repo:app.bsky.feed.post' },
      body: PNG,
    });

    expect(response.status).toBe(403);
    expect((await response.json()).message).toMatch(/image\/png/);
    expect(store.size).toBe(0);
  });
});

// An app-password session names no blob scope and must still be able to upload.
describe('app-password uploads', () => {
  it('accepts a blob from an app-password session', async () => {
    const { handlers, store } = createBlobs();
    const response = await call(handlers, UPLOAD, {
      auth: { did: DID, scope: 'com.atproto.appPass' },
      body: PNG,
    });

    expect(response.status).toBe(200);
    expect(store.size).toBe(1);
  });
});

/**
 * The upload handler off a hand-built request, so the body can be a bare
 * ReadableStream (chunk types and pull timing under test control) or absent,
 * and content-length can be set freely.
 * @param {ReturnType<typeof createBlobHandlers>} handlers
 * @param {{body?: ReadableStream|null, headers?: Record<string, string>, auth?: {did: string, scope: string}}} options
 */
async function upload(handlers, { body = null, headers = {}, auth = FULL }) {
  const route = handlers.routes[UPLOAD];
  const request = /** @type {Request} */ (
    /** @type {unknown} */ ({ headers: new Headers(headers), body })
  );
  const handler =
    /** @type {(request: Request, url: URL, auth: unknown) => Promise<Response>} */ (
      route.handler
    );
  return handler(request, new URL(`https://pds.example.com${UPLOAD}`), auth);
}

/**
 * @param {Array<Uint8Array|ArrayBuffer>} chunks
 * @returns {ReadableStream}
 */
function streamOf(chunks) {
  return new ReadableStream({
    start(controller) {
      for (const chunk of chunks) controller.enqueue(chunk);
      controller.close();
    },
  });
}

describe('streamed uploads', () => {
  it('streams the chunks through the writer', async () => {
    const { handlers, writerLog, meta, store } = createBlobs({
      streaming: true,
    });
    const response = await upload(handlers, {
      body: streamOf([PNG.slice(0, 4), PNG.slice(4, 8), PNG.slice(8)]),
    });
    const body = await response.json();

    expect(response.status).toBe(200);
    expect(body.blob.mimeType).toBe('image/png');
    expect(body.blob.size).toBe(PNG.length);
    expect(body.blob.ref.$link).toBe(cidToString(await createBlobCid(PNG)));
    expect(writerLog.writes).toEqual([4, 4, 2]);
    expect(writerLog.committed).toEqual({
      cid: body.blob.ref.$link,
      mimeType: 'image/png',
    });
    expect(meta.get(body.blob.ref.$link)?.size).toBe(PNG.length);
    expect(store.get(body.blob.ref.$link)?.data).toEqual(PNG);
  });

  it('accepts a chunk that is not a Uint8Array', async () => {
    const { handlers, store } = createBlobs({ streaming: true });
    const response = await upload(handlers, {
      body: streamOf([PNG.slice(0, 6).buffer, PNG.slice(6)]),
    });
    const body = await response.json();

    expect(response.status).toBe(200);
    expect(store.get(body.blob.ref.$link)?.data).toEqual(PNG);
  });

  it('stops a stream that outgrows the limit', async () => {
    const { handlers, writerLog, store } = createBlobs({
      streaming: true,
      blobUploadLimit: 4,
    });
    const response = await upload(handlers, {
      body: streamOf([PNG.slice(0, 3), PNG.slice(3, 6)]),
    });

    expect(response.status).toBe(413);
    expect((await response.json()).error).toBe('PayloadTooLarge');
    expect(writerLog.aborted).toBe(true);
    expect(writerLog.committed).toBe(null);
    expect(store.size).toBe(0);
  });

  it('refuses a declared length over the limit before reading', async () => {
    const { handlers, writerLog } = createBlobs({
      streaming: true,
      blobUploadLimit: 4,
    });
    const response = await upload(handlers, {
      body: streamOf([PNG]),
      headers: { 'content-length': '100' },
    });

    expect(response.status).toBe(413);
    expect(writerLog.writes).toEqual([]);
    expect(writerLog.aborted).toBe(false);
  });

  it('aborts the writer when the scope refuses the type', async () => {
    const { handlers, writerLog, store } = createBlobs({ streaming: true });
    const response = await upload(handlers, {
      body: streamOf([PNG]),
      auth: { did: DID, scope: 'repo:app.bsky.feed.post' },
    });

    expect(response.status).toBe(403);
    expect(writerLog.aborted).toBe(true);
    expect(store.size).toBe(0);
  });

  it('accepts an empty body', async () => {
    const { handlers, meta } = createBlobs({ streaming: true });
    const response = await upload(handlers, {
      headers: { 'content-type': 'application/x-thing' },
    });
    const body = await response.json();

    expect(response.status).toBe(200);
    expect(body.blob.size).toBe(0);
    expect(body.blob.mimeType).toBe('application/x-thing');
    expect(meta.get(body.blob.ref.$link)?.size).toBe(0);
  });

  it('abandons the staged bytes when a write fails', async () => {
    const { handlers, blobs } = createBlobs();
    /** @type {any} */ (blobs).createWriter = async () => ({
      async write() {
        throw new Error('the staging area is full');
      },
      async commit() {},
      async abort() {
        throw new Error('and the abort failed too');
      },
    });

    await expect(upload(handlers, { body: streamOf([PNG]) })).rejects.toThrow(
      /staging area is full/,
    );
  });

  it('hashes with the injected hasher', async () => {
    let used = 0;
    const { handlers } = createBlobs({
      streaming: true,
      hasher: () => {
        used += 1;
        return createSha256();
      },
    });
    const response = await upload(handlers, { body: streamOf([PNG]) });
    const body = await response.json();

    expect(used).toBe(1);
    expect(body.blob.ref.$link).toBe(cidToString(await createBlobCid(PNG)));
  });
});

describe('getBlob', () => {
  it('serves the bytes with headers that stop execution', async () => {
    const { handlers } = createBlobs();
    const cid = (
      await (await call(handlers, UPLOAD, { auth: FULL, body: PNG })).json()
    ).blob.ref.$link;

    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`);

    expect(response.status).toBe(200);
    expect(response.headers.get('Content-Type')).toBe('image/png');
    expect(response.headers.get('X-Content-Type-Options')).toBe('nosniff');
    expect(response.headers.get('Content-Disposition')).toBe(
      `attachment; filename="${cid}"`,
    );
    expect(response.headers.get('Content-Security-Policy')).toBe(
      "default-src 'none'; sandbox",
    );
    expect(new Uint8Array(await response.arrayBuffer())).toEqual(PNG);
  });

  it('needs both did and cid', async () => {
    const { handlers } = createBlobs();
    const response = await call(handlers, `${GET}?did=${DID}`);

    expect(response.status).toBe(400);
    expect((await response.json()).message).toMatch(/Missing required/);
  });

  it('refuses a blob only space records reference', async () => {
    const { handlers, actorStorage } = createBlobs();
    const cid = (
      await (await call(handlers, UPLOAD, { auth: FULL, body: PNG })).json()
    ).blob.ref.$link;
    actorStorage.links.set(cid, [
      `at://did:plc:auth/space/com.example.forum/general/${DID}/com.example.post/a`,
    ]);

    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`);

    expect(response.status).toBe(404);
    expect((await response.json()).error).toBe('BlobNotFound');
  });

  it('serves a blob a public record also references', async () => {
    const { handlers, actorStorage } = createBlobs();
    const cid = (
      await (await call(handlers, UPLOAD, { auth: FULL, body: PNG })).json()
    ).blob.ref.$link;
    actorStorage.links.set(cid, [
      `at://did:plc:auth/space/com.example.forum/general/${DID}/com.example.post/a`,
      `at://${DID}/com.example.post/b`,
    ]);

    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`);
    expect(response.status).toBe(200);
  });

  it('treats a link that is not an at uri as public', async () => {
    const { handlers, actorStorage } = createBlobs();
    const cid = (
      await (await call(handlers, UPLOAD, { auth: FULL, body: PNG })).json()
    ).blob.ref.$link;
    actorStorage.links.set(cid, ['not-a-record-uri']);

    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`);
    expect(response.status).toBe(200);
  });

  it('refuses a did this server does not host', async () => {
    const { handlers } = createBlobs();
    const response = await call(
      handlers,
      `${GET}?did=did:plc:someoneelse&cid=bafkreiabc123def`,
    );

    expect(response.status).toBe(400);
    expect((await response.json()).message).toMatch(/does not match/);
  });

  it('refuses a cid that is not one', async () => {
    const { handlers } = createBlobs();
    const response = await call(handlers, `${GET}?did=${DID}&cid=short`);

    expect(response.status).toBe(400);
    expect((await response.json()).message).toMatch(/Invalid CID/);
  });

  it('is a BlobNotFound for a cid the store does not hold', async () => {
    const { handlers } = createBlobs();
    const response = await call(
      handlers,
      `${GET}?did=${DID}&cid=bafkreiabsentblob`,
    );

    expect(response.status).toBe(404);
    expect((await response.json()).error).toBe('BlobNotFound');
  });
});

describe('getBlob ranges', () => {
  /**
   * @param {{streaming?: boolean}} [options]
   */
  async function withPng({ streaming = false } = {}) {
    const made = createBlobs({ streaming });
    const uploaded = await (
      await call(made.handlers, UPLOAD, { auth: FULL, body: PNG })
    ).json();
    return { ...made, cid: uploaded.blob.ref.$link };
  }

  it('answers the whole blob with length and range advertising', async () => {
    const { handlers, cid } = await withPng();
    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`);

    expect(response.status).toBe(200);
    expect(response.headers.get('Accept-Ranges')).toBe('bytes');
    expect(response.headers.get('Content-Length')).toBe(String(PNG.length));
    expect(response.headers.get('Content-Range')).toBe(null);
  });

  it('slices a range from a buffered port', async () => {
    const { handlers, cid } = await withPng();
    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`, {
      headers: { range: 'bytes=2-5' },
    });

    expect(response.status).toBe(206);
    expect(response.headers.get('Content-Range')).toBe(
      `bytes 2-5/${PNG.length}`,
    );
    expect(response.headers.get('Content-Length')).toBe('4');
    expect(new Uint8Array(await response.arrayBuffer())).toEqual(
      PNG.subarray(2, 6),
    );
  });

  it('passes a clamped range to the streaming port', async () => {
    const { handlers, cid, streamLog } = await withPng({ streaming: true });
    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`, {
      headers: { range: 'bytes=4-' },
    });

    expect(response.status).toBe(206);
    expect(streamLog).toEqual([{ start: 4, end: PNG.length - 1 }]);
    expect(response.headers.get('Content-Range')).toBe(
      `bytes 4-${PNG.length - 1}/${PNG.length}`,
    );
    expect(new Uint8Array(await response.arrayBuffer())).toEqual(
      PNG.subarray(4),
    );
  });

  it('streams the whole blob when no range is asked', async () => {
    const { handlers, cid, streamLog } = await withPng({ streaming: true });
    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`);

    expect(response.status).toBe(200);
    expect(streamLog).toEqual([null]);
    expect(new Uint8Array(await response.arrayBuffer())).toEqual(PNG);
  });

  it('serves a suffix range', async () => {
    const { handlers, cid } = await withPng();
    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`, {
      headers: { range: 'bytes=-3' },
    });

    expect(response.status).toBe(206);
    expect(new Uint8Array(await response.arrayBuffer())).toEqual(
      PNG.subarray(PNG.length - 3),
    );
  });

  it('refuses a range past the end', async () => {
    const { handlers, cid } = await withPng();
    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`, {
      headers: { range: `bytes=${PNG.length}-` },
    });

    expect(response.status).toBe(416);
    expect(response.headers.get('Content-Range')).toBe(`bytes */${PNG.length}`);
  });

  it('ignores a malformed range header', async () => {
    const { handlers, cid } = await withPng();
    const response = await call(handlers, `${GET}?did=${DID}&cid=${cid}`, {
      headers: { range: 'rows=1-2' },
    });

    expect(response.status).toBe(200);
  });

  it('serves bytes that have no metadata row, whole and unranged', async () => {
    const { handlers, store } = createBlobs();
    store.set('bafkreinometa', { data: PNG, mimeType: 'image/png' });

    const response = await call(
      handlers,
      `${GET}?did=${DID}&cid=bafkreinometa`,
      { headers: { range: 'bytes=0-3' } },
    );

    expect(response.status).toBe(200);
    expect(response.headers.get('Content-Length')).toBe(String(PNG.length));
    expect(new Uint8Array(await response.arrayBuffer())).toEqual(PNG);
  });

  it('is a BlobNotFound when the stream port holds nothing', async () => {
    const { handlers, meta } = createBlobs({ streaming: true });
    meta.set('bafkreimetaonly', { mimeType: 'image/png', size: 5 });

    const response = await call(
      handlers,
      `${GET}?did=${DID}&cid=bafkreimetaonly`,
    );

    expect(response.status).toBe(404);
  });

  it('is a BlobNotFound when the buffered port holds nothing', async () => {
    const { handlers, meta } = createBlobs();
    meta.set('bafkreimetaonly', { mimeType: 'image/png', size: 5 });

    const response = await call(
      handlers,
      `${GET}?did=${DID}&cid=bafkreimetaonly`,
    );

    expect(response.status).toBe(404);
  });
});

describe('listBlobs', () => {
  it('lists the stored cids', async () => {
    const { handlers, meta } = createBlobs();
    meta.set('bafkreib', { mimeType: 'image/png', size: 1 });
    meta.set('bafkreia', { mimeType: 'image/png', size: 1 });

    const body = await (await call(handlers, LIST)).json();

    expect(body.cids).toEqual(['bafkreia', 'bafkreib']);
  });

  it('leaves out a blob only space records reference', async () => {
    const { handlers, meta, actorStorage } = createBlobs();
    meta.set('bafkreia', { mimeType: 'image/png', size: 1 });
    meta.set('bafkreib', { mimeType: 'image/png', size: 1 });
    actorStorage.links.set('bafkreib', [
      `at://did:plc:auth/space/com.example.forum/general/${DID}/com.example.post/a`,
    ]);

    const body = await (await call(handlers, LIST)).json();

    expect(body.cids).toEqual(['bafkreia']);
  });

  it('omits the cursor on a final page', async () => {
    const { handlers, meta } = createBlobs();
    meta.set('bafkreia', { mimeType: 'image/png', size: 1 });

    const body = await (await call(handlers, LIST)).json();

    expect(body.cids).toEqual(['bafkreia']);
    expect('cursor' in body).toBe(false);
  });

  it('gives a cursor while more pages remain', async () => {
    const { handlers, meta } = createBlobs();
    for (const c of ['bafkreia', 'bafkreib', 'bafkreic']) {
      meta.set(c, { mimeType: 'image/png', size: 1 });
    }

    const first = await (await call(handlers, `${LIST}?limit=2`)).json();
    expect(first.cids).toEqual(['bafkreia', 'bafkreib']);
    expect(first.cursor).toBe('bafkreib');

    const second = await (
      await call(handlers, `${LIST}?limit=2&cursor=${first.cursor}`)
    ).json();
    expect(second.cids).toEqual(['bafkreic']);
    expect('cursor' in second).toBe(false);
  });

  it('holds the page size to 1000', async () => {
    const { handlers, actorStorage } = createBlobs();
    /** @type {number[]} */
    const limits = [];
    const original = actorStorage.listBlobs;
    actorStorage.listBlobs = async (cursor, limit) => {
      limits.push(limit);
      return original(cursor, limit);
    };

    await call(handlers, `${LIST}?limit=5000`);
    expect(limits).toEqual([1000]);
  });
});