import { createHash } from "node:crypto";
import type { MediaUpload } from "@rakazo/contracts";
import { MEDIA_UPLOAD_CHUNK_BYTES } from "@rakazo/contracts";
import { describe, expect, it, vi } from "vitest";
import type { ChunkBytes, UploadTransport } from "./media-upload";
import {
  chunkCount,
  chunkRange,
  combinedProgress,
  hashOriginal,
  isRetryableUploadError,
  isUploadExpired,
  MediaUploadHttpError,
  missingChunks,
  retryDelayMs,
  uploadOriginal,
} from "./media-upload";

function bytes(size: number): ChunkBytes {
  const out = new Uint8Array(size);
  for (let index = 0; index < size; index++) out[index] = (index * 31 + 7) & 0xff;
  return out;
}

const sha = (data: Uint8Array) => createHash("sha256").update(data).digest("hex");

/** In-memory server with the same chunk rules as the API route. */
function fakeServer(file: Uint8Array, chunkBytes: number) {
  let nextId = 0;
  const uploads = new Map<string, { upload: MediaUpload; chunks: Map<number, Uint8Array> }>();
  const failures: Array<(index: number) => Error | null> = [];
  const view = (id: string): MediaUpload => {
    const entry = uploads.get(id);
    if (!entry) throw new MediaUploadHttpError(404, "Not found");
    return { ...entry.upload, receivedChunks: [...entry.chunks.keys()].sort((a, b) => a - b) };
  };
  const transport = {
    begin: vi.fn(async (input: { assetId: string; sizeBytes: number; sha256: string }) => {
      const open = [...uploads.values()].find(
        (entry) => entry.upload.assetId === input.assetId && entry.upload.status === "open",
      );
      if (open) return view(open.upload.id);
      nextId += 1;
      const id = `up${nextId}`;
      uploads.set(id, {
        upload: {
          id,
          assetId: input.assetId,
          sizeBytes: input.sizeBytes,
          chunkBytes,
          receivedChunks: [],
          status: "open",
        },
        chunks: new Map(),
      });
      return view(id);
    }),
    status: vi.fn(async (uploadId: string) => view(uploadId)),
    putChunk: vi.fn(async (uploadId: string, index: number, chunk: Uint8Array) => {
      const failure = failures.shift()?.(index);
      if (failure) throw failure;
      const entry = uploads.get(uploadId);
      if (!entry) throw new MediaUploadHttpError(404, "Not found");
      if (entry.upload.status !== "open") throw new MediaUploadHttpError(409, "Upload expired");
      const { length } = chunkRange(index, entry.upload.sizeBytes, chunkBytes);
      if (chunk.byteLength !== length) throw new MediaUploadHttpError(400, "Chunk is incomplete");
      entry.chunks.set(index, chunk.slice());
      return view(uploadId);
    }),
    complete: vi.fn(async (uploadId: string) => {
      const entry = uploads.get(uploadId);
      if (!entry) throw new MediaUploadHttpError(404, "Not found");
      if (entry.upload.status !== "open") throw new MediaUploadHttpError(409, "Upload expired");
      const count = chunkCount(entry.upload.sizeBytes, chunkBytes);
      const parts: Uint8Array[] = [];
      for (let index = 0; index < count; index++) {
        const part = entry.chunks.get(index);
        if (!part) throw new MediaUploadHttpError(409, "Upload is missing chunks");
        parts.push(part);
      }
      const assembled = new Uint8Array(entry.upload.sizeBytes);
      let offset = 0;
      for (const part of parts) {
        assembled.set(part, offset);
        offset += part.byteLength;
      }
      expect(sha(assembled)).toBe(sha(file));
      entry.upload.status = "complete";
      return view(uploadId);
    }),
  } satisfies UploadTransport;
  return {
    transport,
    uploads,
    failNext: (fail: (index: number) => Error | null) => failures.push(fail),
    expire: (id: string) => {
      const entry = uploads.get(id);
      if (entry) entry.upload.status = "failed";
    },
  };
}

const read = (file: ChunkBytes) => (start: number, length: number) =>
  file.slice(start, start + length);
const noWait = async () => undefined;

describe("chunk planning", () => {
  it("splits a file into full chunks and a short last chunk", () => {
    expect(chunkCount(10, 4)).toBe(3);
    expect(chunkRange(0, 10, 4)).toEqual({ start: 0, length: 4 });
    expect(chunkRange(1, 10, 4)).toEqual({ start: 4, length: 4 });
    expect(chunkRange(2, 10, 4)).toEqual({ start: 8, length: 2 });
  });

  it("uses exact chunks when the size divides evenly", () => {
    expect(chunkCount(8, 4)).toBe(2);
    expect(chunkRange(1, 8, 4)).toEqual({ start: 4, length: 4 });
  });

  it("matches the contract chunk size for large originals", () => {
    const size = 3 * MEDIA_UPLOAD_CHUNK_BYTES + 5;
    expect(chunkCount(size, MEDIA_UPLOAD_CHUNK_BYTES)).toBe(4);
    expect(chunkRange(3, size, MEDIA_UPLOAD_CHUNK_BYTES)).toEqual({
      start: 3 * MEDIA_UPLOAD_CHUNK_BYTES,
      length: 5,
    });
  });

  it("rejects invalid sizes and indexes", () => {
    expect(() => chunkCount(0, 4)).toThrow();
    expect(() => chunkCount(10, 0)).toThrow();
    expect(() => chunkRange(3, 10, 4)).toThrow();
    expect(() => chunkRange(-1, 10, 4)).toThrow();
  });

  it("lists the chunks the server has not confirmed", () => {
    expect(missingChunks({ sizeBytes: 10, chunkBytes: 4, receivedChunks: [] })).toEqual([0, 1, 2]);
    expect(missingChunks({ sizeBytes: 10, chunkBytes: 4, receivedChunks: [2, 0] })).toEqual([1]);
    expect(missingChunks({ sizeBytes: 10, chunkBytes: 4, receivedChunks: [0, 1, 2] })).toEqual([]);
  });
});

describe("error classification", () => {
  it("recognizes an expired upload", () => {
    expect(isUploadExpired(new MediaUploadHttpError(409, "Upload expired"))).toBe(true);
    expect(isUploadExpired(new MediaUploadHttpError(409, "Upload is missing chunks"))).toBe(false);
    expect(isUploadExpired(new Error("Upload expired"))).toBe(false);
  });

  it("retries network and server failures, not client errors", () => {
    expect(isRetryableUploadError(new TypeError("Network request failed"))).toBe(true);
    expect(isRetryableUploadError(new MediaUploadHttpError(503, "busy"))).toBe(true);
    expect(isRetryableUploadError(new MediaUploadHttpError(408, "slow"))).toBe(true);
    expect(isRetryableUploadError(new MediaUploadHttpError(400, "Chunk is incomplete"))).toBe(
      false,
    );
    expect(isRetryableUploadError(new MediaUploadHttpError(413, "Chunk is too large"))).toBe(false);
  });

  it("backs off exponentially up to a cap", () => {
    expect(retryDelayMs(1)).toBe(1_000);
    expect(retryDelayMs(3)).toBe(4_000);
    expect(retryDelayMs(20)).toBe(30_000);
  });
});

describe("hashOriginal", () => {
  it("hashes the file in pieces", async () => {
    const file = bytes(1_000);
    const progress: number[] = [];
    const hash = await hashOriginal(file.byteLength, read(file), {
      readBytes: 64,
      onProgress: (fraction) => progress.push(fraction),
    });
    expect(hash).toBe(sha(file));
    expect(progress.at(-1)).toBe(1);
  });

  it("fails when the file is shorter than reported", async () => {
    const file = bytes(100);
    await expect(hashOriginal(200, read(file), { readBytes: 64 })).rejects.toThrow(/changed/);
  });
});

describe("uploadOriginal", () => {
  const file = bytes(10_000);
  const chunkBytes = 4_096;
  const base = {
    assetId: "asset1",
    fileName: "clip.mov",
    sizeBytes: file.byteLength,
    sha256: sha(file),
  };

  it("uploads every chunk in order and completes", async () => {
    const server = fakeServer(file, chunkBytes);
    const progress: number[] = [];
    const upload = await uploadOriginal({
      ...base,
      read: read(file),
      transport: server.transport,
      onProgress: (fraction) => progress.push(fraction),
      wait: noWait,
    });
    expect(upload.status).toBe("complete");
    expect(server.transport.putChunk.mock.calls.map(([, index]) => index)).toEqual([0, 1, 2]);
    expect(server.transport.complete).toHaveBeenCalledTimes(1);
    expect(progress.at(-1)).toBe(1);
  });

  it("resumes from receivedChunks of an open upload", async () => {
    const server = fakeServer(file, chunkBytes);
    const first = await server.transport.begin(base);
    await server.transport.putChunk(first.id, 0, read(file)(0, chunkBytes));
    await server.transport.putChunk(
      first.id,
      2,
      read(file)(2 * chunkBytes, 10_000 - 2 * chunkBytes),
    );
    server.transport.putChunk.mockClear();

    const upload = await uploadOriginal({
      ...base,
      read: read(file),
      transport: server.transport,
      wait: noWait,
    });
    expect(upload.id).toBe(first.id);
    expect(server.transport.putChunk.mock.calls.map(([, index]) => index)).toEqual([1]);
  });

  it("asks the server what arrived after a network failure and continues", async () => {
    const server = fakeServer(file, chunkBytes);
    server.failNext(() => null);
    server.failNext(() => new TypeError("Network request failed"));
    const wait = vi.fn(noWait);
    const upload = await uploadOriginal({
      ...base,
      read: read(file),
      transport: server.transport,
      wait,
    });
    expect(upload.status).toBe("complete");
    expect(server.transport.status).toHaveBeenCalledTimes(1);
    expect(wait).toHaveBeenCalledWith(1_000, undefined);
    // Chunk 0 landed before the failure; only 1 is sent again.
    expect(server.transport.putChunk.mock.calls.map(([, index]) => index)).toEqual([0, 1, 1, 2]);
  });

  it("begins again when the upload expired", async () => {
    const server = fakeServer(file, chunkBytes);
    server.failNext(() => null);
    server.failNext(() => {
      server.expire("up1");
      return new MediaUploadHttpError(409, "Upload expired");
    });
    const upload = await uploadOriginal({
      ...base,
      read: read(file),
      transport: server.transport,
      wait: noWait,
    });
    expect(upload.status).toBe("complete");
    expect(upload.id).toBe("up2");
    expect(server.transport.begin).toHaveBeenCalledTimes(2);
    // The new upload starts from chunk 0.
    expect(server.transport.putChunk.mock.calls.map(([id, index]) => `${id}:${index}`)).toEqual([
      "up1:0",
      "up1:1",
      "up2:0",
      "up2:1",
      "up2:2",
    ]);
  });

  it("begins again when completion reports the upload expired", async () => {
    const server = fakeServer(file, chunkBytes);
    const complete = server.transport.complete.getMockImplementation();
    server.transport.complete.mockImplementationOnce(async (uploadId: string) => {
      server.expire(uploadId);
      throw new MediaUploadHttpError(409, "Upload expired");
    });
    if (complete) server.transport.complete.mockImplementation(complete);
    const upload = await uploadOriginal({
      ...base,
      read: read(file),
      transport: server.transport,
      wait: noWait,
    });
    expect(upload.status).toBe("complete");
    expect(server.transport.begin).toHaveBeenCalledTimes(2);
  });

  it("stops restarting after the restart budget", async () => {
    const server = fakeServer(file, chunkBytes);
    for (let attempt = 0; attempt < 5; attempt++) {
      server.failNext(() => {
        for (const id of server.uploads.keys()) server.expire(id);
        return new MediaUploadHttpError(409, "Upload expired");
      });
    }
    await expect(
      uploadOriginal({
        ...base,
        read: read(file),
        transport: server.transport,
        wait: noWait,
        maxRestarts: 2,
      }),
    ).rejects.toThrow("Upload expired");
    expect(server.transport.begin).toHaveBeenCalledTimes(3);
  });

  it("gives up after repeated failures", async () => {
    const server = fakeServer(file, chunkBytes);
    for (let attempt = 0; attempt < 10; attempt++) {
      server.failNext(() => new MediaUploadHttpError(503, "busy"));
    }
    await expect(
      uploadOriginal({
        ...base,
        read: read(file),
        transport: server.transport,
        wait: noWait,
        maxAttempts: 3,
      }),
    ).rejects.toThrow("busy");
    expect(server.transport.putChunk).toHaveBeenCalledTimes(3);
  });

  it("does not retry a rejected chunk", async () => {
    const server = fakeServer(file, chunkBytes);
    server.failNext(() => new MediaUploadHttpError(413, "Chunk is too large"));
    await expect(
      uploadOriginal({ ...base, read: read(file), transport: server.transport, wait: noWait }),
    ).rejects.toThrow("Chunk is too large");
    expect(server.transport.putChunk).toHaveBeenCalledTimes(1);
  });

  it("returns at once when the upload is already complete", async () => {
    const server = fakeServer(file, chunkBytes);
    await uploadOriginal({ ...base, read: read(file), transport: server.transport, wait: noWait });
    server.transport.putChunk.mockClear();
    server.transport.begin.mockImplementationOnce(async () => ({
      ...(await server.transport.status("up1")),
    }));
    const again = await uploadOriginal({
      ...base,
      read: read(file),
      transport: server.transport,
      wait: noWait,
    });
    expect(again.status).toBe("complete");
    expect(server.transport.putChunk).not.toHaveBeenCalled();
  });

  it("stops when aborted", async () => {
    const server = fakeServer(file, chunkBytes);
    const controller = new AbortController();
    server.failNext(() => {
      controller.abort(new Error("stopped"));
      return new TypeError("Network request failed");
    });
    await expect(
      uploadOriginal({
        ...base,
        read: read(file),
        transport: server.transport,
        wait: noWait,
        signal: controller.signal,
      }),
    ).rejects.toThrow("stopped");
    expect(server.transport.status).not.toHaveBeenCalled();
  });

  it("fails when the file changed size while uploading", async () => {
    const server = fakeServer(file, chunkBytes);
    const short = (start: number, length: number) => file.slice(start, start + length - 1);
    await expect(
      uploadOriginal({ ...base, read: short, transport: server.transport, wait: noWait }),
    ).rejects.toThrow(/changed/);
    expect(server.transport.status).not.toHaveBeenCalled();
  });
});

describe("combinedProgress", () => {
  it("weights uploads by size", () => {
    expect(combinedProgress([])).toBe(0);
    expect(
      combinedProgress([
        { sizeBytes: 100, fraction: 1 },
        { sizeBytes: 300, fraction: 0 },
      ]),
    ).toBe(0.25);
  });
});