import { createHash } from "node:crypto";
import fs from "node:fs/promises";
import path from "node:path";
import { afterEach, describe, expect, test } from "vitest";
import {
  downloadTelegramImage,
  imageAttachmentMetadata,
  isImageDocument,
  resolveImageArtifact,
  selectPhoto,
  TELEGRAM_IMAGE_MAX_BYTES,
  type TelegramPhotoSize,
  type TelegramFile,
} from "../src/connectors/telegram-images.js";
import { TELEGRAM_IMAGE_MAX_BASE64_CHARS } from "../src/connectors/telegram-image-contract.js";
import { SANDBOX_PROTOCOL_VERSION, sandboxRunPacketSchema } from "../src/agents/sandbox/protocol.js";
import { temporaryProject } from "./helpers.js";

const roots: string[] = [];

afterEach(async () => {
  await Promise.all(roots.splice(0).map((root) => fs.rm(root, { recursive: true, force: true })));
});

// Minimal 1x1 PNG image (67 bytes)
const PNG_BYTES = Buffer.from([
  0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,
  0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52,
  0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
  0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, 0x15, 0xc4,
  0x89, 0x00, 0x00, 0x00, 0x0d, 0x49, 0x44, 0x41,
  0x54, 0x78, 0x9c, 0x62, 0x00, 0x01, 0x00, 0x00,
  0x05, 0x00, 0x01, 0x0d, 0x0a, 0x2d, 0xb4, 0x00,
  0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae,
  0x42, 0x60, 0x82,
]);

// Minimal JPEG header + footer (not a valid JPEG but sufficient for magic byte tests)
const JPEG_BYTES = Buffer.concat([
  Buffer.from([0xff, 0xd8, 0xff, 0xe0, 0x00, 0x10, 0x4a, 0x46, 0x49, 0x46]),
  Buffer.alloc(100, 0x00),
  Buffer.from([0xff, 0xd9]),
]);

const INVALID_BYTES = Buffer.from([0x00, 0x01, 0x02, 0x03]);

describe("selectPhoto", () => {
  test("returns the largest photo within the byte limit", () => {
    const photos: TelegramPhotoSize[] = [
      { file_id: "small", file_unique_id: "u1", width: 100, height: 100, file_size: 10_000 },
      { file_id: "large", file_unique_id: "u2", width: 800, height: 800, file_size: 5_000_000 },
    ];
    const selected = selectPhoto(photos);
    expect(selected).toBeDefined();
    expect(selected!.file_id).toBe("large");
  });

  test("returns undefined when all photos exceed the byte limit", () => {
    const photos: TelegramPhotoSize[] = [
      { file_id: "big", file_unique_id: "u1", width: 800, height: 800, file_size: TELEGRAM_IMAGE_MAX_BYTES + 1 },
    ];
    const selected = selectPhoto(photos);
    expect(selected).toBeUndefined();
  });

  test("returns undefined for empty array", () => {
    expect(selectPhoto([])).toBeUndefined();
  });
});

describe("isImageDocument", () => {
  test("returns true for PNG documents", () => {
    const file: TelegramFile = { file_id: "f1", file_unique_id: "u1", mime_type: "image/png" };
    expect(isImageDocument(file)).toBe(true);
  });

  test("returns true for JPEG documents", () => {
    const file: TelegramFile = { file_id: "f1", file_unique_id: "u1", mime_type: "image/jpeg" };
    expect(isImageDocument(file)).toBe(true);
  });

  test("returns false for non-image documents", () => {
    const file: TelegramFile = { file_id: "f1", file_unique_id: "u1", mime_type: "application/pdf" };
    expect(isImageDocument(file)).toBe(false);
  });

  test("returns false for documents without MIME type", () => {
    const file: TelegramFile = { file_id: "f1", file_unique_id: "u1" };
    expect(isImageDocument(file)).toBe(false);
  });
});

describe("downloadTelegramImage", () => {
  test("downloads, validates, and stores a PNG image content-addressed", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const expectedHash = createHash("sha256").update(PNG_BYTES).digest("hex");

    const fetchImpl = mockFetchImpl({
      getFile: { file_path: "photos/file_123.png" },
      file: PNG_BYTES,
    });

    const artifact = await downloadTelegramImage("file-id-123", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    });

    expect(artifact.sha256).toBe(expectedHash);
    expect(artifact.mimeType).toBe("image/png");
    expect(artifact.sizeBytes).toBe(PNG_BYTES.byteLength);
    expect(artifact.relativePath).toBe(`sha256/${expectedHash.slice(0, 2)}/${expectedHash}`);

    // Verify the file was written
    const written = await fs.readFile(path.join(artifactRoot, artifact.relativePath));
    expect(written.equals(PNG_BYTES)).toBe(true);
  });

  test("downloads, validates, and stores a JPEG image content-addressed", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const expectedHash = createHash("sha256").update(JPEG_BYTES).digest("hex");

    const fetchImpl = mockFetchImpl({
      getFile: { file_path: "photos/file_456.jpg" },
      file: JPEG_BYTES,
    });

    const artifact = await downloadTelegramImage("file-id-456", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    });

    expect(artifact.sha256).toBe(expectedHash);
    expect(artifact.mimeType).toBe("image/jpeg");
    expect(artifact.sizeBytes).toBe(JPEG_BYTES.byteLength);
  });

  test("rejects JPEG-looking bytes without the terminal EOI marker", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });
    const truncatedJpeg = JPEG_BYTES.subarray(0, -2);
    const fetchImpl = mockFetchImpl({ getFile: { file_path: "photos/truncated.jpg" }, file: truncatedJpeg });
    await expect(downloadTelegramImage("file-truncated", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow("unsupported or unrecognized format");
  });

  test("aligns the raw image cap with the sandbox base64 character cap", () => {
    const data = Buffer.alloc(TELEGRAM_IMAGE_MAX_BYTES).toString("base64");
    expect(data.length).toBe(TELEGRAM_IMAGE_MAX_BASE64_CHARS);
    const packet = {
      version: SANDBOX_PROTOCOL_VERSION,
      runId: "run-image-bound",
      systemPrompt: "system",
      prompt: "prompt",
      images: [{ type: "image", data, mimeType: "image/png" }],
      model: {
        provider: "tinker",
        id: "thinkingmachines/Inkling-Small",
        reasoning: true,
        acceptsImages: true,
        jsonObjectResponseFormat: false,
        contextWindow: 32_000,
        maxTokens: 1_000,
      },
      broker: {
        socketPath: "/broker/provider.sock",
        capability: "a".repeat(64),
        origin: "http://provider.invalid",
        route: "/v1/chat/completions",
      },
    };
    expect(sandboxRunPacketSchema.safeParse(packet).success).toBe(true);
    expect(sandboxRunPacketSchema.safeParse({
      ...packet,
      images: [{ ...packet.images[0], data: `${data}A` }],
    }).success).toBe(false);
  });

  test("rejects unsupported image formats", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const fetchImpl = mockFetchImpl({
      getFile: { file_path: "photos/file_bad.bin" },
      file: INVALID_BYTES,
    });

    await expect(downloadTelegramImage("file-bad", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow("unsupported or unrecognized format");
  });

  test("rejects empty files", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const fetchImpl = mockFetchImpl({
      getFile: { file_path: "photos/empty.png" },
      file: Buffer.alloc(0),
    });

    await expect(downloadTelegramImage("file-empty", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow("empty");
  });

  test("rejects oversize files", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const oversize = Buffer.concat([PNG_BYTES, Buffer.alloc(TELEGRAM_IMAGE_MAX_BYTES + 1, 0x00)]);

    const fetchImpl = mockFetchImpl({
      getFile: { file_path: "photos/huge.png" },
      file: oversize,
    });

    await expect(downloadTelegramImage("file-huge", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow("exceeds");
  });

  test("rejects an announced oversize body before reading it", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });
    const fetchImpl = (async (input: RequestInfo | URL) => {
      const url = String(input);
      if (url.includes("/getFile")) {
        return new Response(JSON.stringify({ ok: true, result: { file_path: "photos/huge.png" } }), {
          status: 200,
          headers: { "content-type": "application/json" },
        });
      }
      return new Response(new Uint8Array(PNG_BYTES), {
        status: 200,
        headers: { "content-length": String(TELEGRAM_IMAGE_MAX_BYTES + 1) },
      });
    }) as typeof fetch;
    await expect(downloadTelegramImage("file-huge", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow("exceeds");
  });

  test("rejects unsafe getFile paths", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });
    const fetchImpl = mockFetchImpl({ getFile: { file_path: "../token-leak.png" }, file: PNG_BYTES });
    await expect(downloadTelegramImage("file-unsafe", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow("invalid file path");
  });

  test("rejects redirects outside the exact Bot API file boundary", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });
    const fetchImpl = (async (input: RequestInfo | URL) => {
      const url = String(input);
      if (url.includes("/getFile")) {
        return new Response(JSON.stringify({ ok: true, result: { file_path: "photos/image.png" } }), {
          status: 200,
          headers: { "content-type": "application/json" },
        });
      }
      return new Response(null, { status: 302, headers: { location: "https://example.com/image.png" } });
    }) as typeof fetch;
    await expect(downloadTelegramImage("file-redirect", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow("trusted Bot API file boundary");
  });

  test("rejects a symlinked artifact root before writing bytes", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const target = path.join(project, "real-artifacts");
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(target);
    await fs.symlink(target, artifactRoot);
    const fetchImpl = mockFetchImpl({ getFile: { file_path: "photos/image.png" }, file: PNG_BYTES });
    await expect(downloadTelegramImage("file-symlink-root", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow(/real directory|symlinked parents/);
    expect(await fs.readdir(target)).toEqual([]);
  });

  test("rejects symlinked content-addressing parents before writing bytes", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    const escape = path.join(project, "escape");
    await fs.mkdir(artifactRoot);
    await fs.mkdir(escape);
    await fs.symlink(escape, path.join(artifactRoot, "sha256"));
    const fetchImpl = mockFetchImpl({ getFile: { file_path: "photos/image.png" }, file: PNG_BYTES });
    await expect(downloadTelegramImage("file-symlink-parent", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    })).rejects.toThrow(/real directory|symlinked parents/);
    expect(await fs.readdir(escape)).toEqual([]);
  });

  test("is idempotent: downloading the same image twice writes the same file", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const fetchImpl = mockFetchImpl({
      getFile: { file_path: "photos/file_123.png" },
      file: PNG_BYTES,
    });

    const first = await downloadTelegramImage("file-id-123", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    });

    const second = await downloadTelegramImage("file-id-123", {
      baseUrl: "https://api.telegram.org",
      token: "test-token",
      artifactRoot,
      fetchImpl,
    });

    expect(first.relativePath).toBe(second.relativePath);
    expect(first.sha256).toBe(second.sha256);
  });
});

describe("resolveImageArtifact", () => {
  test("resolves a valid artifact to base64 ImageContent", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const hash = createHash("sha256").update(PNG_BYTES).digest("hex");
    const relativePath = `sha256/${hash.slice(0, 2)}/${hash}`;
    const destination = path.join(artifactRoot, relativePath);
    await fs.mkdir(path.dirname(destination), { recursive: true });
    await fs.writeFile(destination, PNG_BYTES);

    const image = await resolveImageArtifact({
      path: relativePath,
      sha256: hash,
      mimeType: "image/png",
      sizeBytes: PNG_BYTES.byteLength,
    }, artifactRoot);

    expect(image.type).toBe("image");
    expect(image.mimeType).toBe("image/png");
    expect(image.data).toBe(PNG_BYTES.toString("base64"));
  });

  test("rejects path traversal attempts", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    await expect(resolveImageArtifact({
      path: "../../../etc/passwd",
      sha256: "fake",
      mimeType: "image/png",
      sizeBytes: 100,
    }, artifactRoot)).rejects.toThrow("path escape");
  });

  test("rejects non-normalized paths", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    await expect(resolveImageArtifact({
      path: "sha256/ab/abc/../../bad",
      sha256: "fake",
      mimeType: "image/png",
      sizeBytes: 100,
    }, artifactRoot)).rejects.toThrow();
  });

  test("rejects paths outside the sha256 directory", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    await expect(resolveImageArtifact({
      path: "other/dir/file.png",
      sha256: "fake",
      mimeType: "image/png",
      sizeBytes: 100,
    }, artifactRoot)).rejects.toThrow("content-addressed directory");
  });

  test("rejects missing files", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    await expect(resolveImageArtifact({
      path: "sha256/ab/abc123",
      sha256: "abc123",
      mimeType: "image/png",
      sizeBytes: 100,
    }, artifactRoot)).rejects.toThrow("missing");
  });

  test("rejects hash mismatch", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const hash = createHash("sha256").update(PNG_BYTES).digest("hex");
    const relativePath = `sha256/${hash.slice(0, 2)}/${hash}`;
    const destination = path.join(artifactRoot, relativePath);
    await fs.mkdir(path.dirname(destination), { recursive: true });
    await fs.writeFile(destination, PNG_BYTES);

    await expect(resolveImageArtifact({
      path: relativePath,
      sha256: "wronghash",
      mimeType: "image/png",
      sizeBytes: PNG_BYTES.byteLength,
    }, artifactRoot)).rejects.toThrow("hash mismatch");
  });

  test("rejects size mismatch", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const hash = createHash("sha256").update(PNG_BYTES).digest("hex");
    const relativePath = `sha256/${hash.slice(0, 2)}/${hash}`;
    const destination = path.join(artifactRoot, relativePath);
    await fs.mkdir(path.dirname(destination), { recursive: true });
    await fs.writeFile(destination, PNG_BYTES);

    await expect(resolveImageArtifact({
      path: relativePath,
      sha256: hash,
      mimeType: "image/png",
      sizeBytes: PNG_BYTES.byteLength + 1,
    }, artifactRoot)).rejects.toThrow("size mismatch");
  });

  test("rejects MIME mismatch", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });

    const hash = createHash("sha256").update(PNG_BYTES).digest("hex");
    const relativePath = `sha256/${hash.slice(0, 2)}/${hash}`;
    const destination = path.join(artifactRoot, relativePath);
    await fs.mkdir(path.dirname(destination), { recursive: true });
    await fs.writeFile(destination, PNG_BYTES);

    await expect(resolveImageArtifact({
      path: relativePath,
      sha256: hash,
      mimeType: "image/jpeg",
      sizeBytes: PNG_BYTES.byteLength,
    }, artifactRoot)).rejects.toThrow("MIME mismatch");
  });

  test("rejects a symlink at the expected content address", async () => {
    const project = await temporaryProject("telegram-image-test-");
    roots.push(project);
    const artifactRoot = path.join(project, "artifacts");
    await fs.mkdir(artifactRoot, { recursive: true });
    const hash = createHash("sha256").update(PNG_BYTES).digest("hex");
    const relativePath = `sha256/${hash.slice(0, 2)}/${hash}`;
    const destination = path.join(artifactRoot, relativePath);
    const target = path.join(project, "target.png");
    await fs.mkdir(path.dirname(destination), { recursive: true });
    await fs.writeFile(target, PNG_BYTES);
    await fs.symlink(target, destination);

    await expect(resolveImageArtifact({
      path: relativePath,
      sha256: hash,
      mimeType: "image/png",
      sizeBytes: PNG_BYTES.byteLength,
    }, artifactRoot)).rejects.toThrow("symlink");
  });
});

describe("imageAttachmentMetadata", () => {
  test("produces bounded metadata with opaque reference only", () => {
    const metadata = imageAttachmentMetadata({
      relativePath: "sha256/ab/abc123",
      sha256: "abc123",
      mimeType: "image/png",
      sizeBytes: 1024,
    });

    expect(metadata).toEqual({
      kind: "image",
      status: "stored",
      mimeType: "image/png",
      sizeBytes: 1024,
      sha256: "abc123",
      artifactPath: "sha256/ab/abc123",
    });

    // Must never contain raw bytes, base64, token, URL, or absolute path
    expect(metadata).not.toHaveProperty("data");
    expect(metadata).not.toHaveProperty("token");
    expect(metadata).not.toHaveProperty("url");
    expect(String(metadata.artifactPath)).not.toContain("..");
    expect(String(metadata.artifactPath)).not.toMatch(/^\//);
  });
});

/**
 * Create a mock fetch implementation that simulates the Telegram Bot API:
 * - POST /bot{token}/getFile returns the file path
 * - GET /file/bot{token}/{file_path} returns the file bytes
 */
function mockFetchImpl(options: {
  getFile: { file_path: string };
  file: Buffer;
}): typeof fetch {
  const { getFile: getFileResponse, file } = options;
  return (async (input: RequestInfo | URL, init?: RequestInit) => {
    const url = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
    const method = init?.method ?? "GET";

    if (url.includes("/getFile")) {
      return new Response(JSON.stringify({
        ok: true,
        result: { file_path: getFileResponse.file_path },
      }), { status: 200, headers: { "content-type": "application/json" } });
    }

    if (url.includes("/file/bot")) {
      return new Response(new Uint8Array(file), { status: 200, headers: { "content-type": "application/octet-stream" } });
    }

    return new Response("not found", { status: 404 });
  }) as typeof fetch;
}