import { encodeSchema } from "@noya-app/noya-schemas";
import { describe, expect, it } from "bun:test";
import { Type } from "@sinclair/typebox";
import { AIManager } from "../AIManager";
import type {
  AIGenerateImageRequest,
  AIGenerateRequest,
  AIGenerateTextRequest,
} from "../rpc/routes";
import { RPCManager } from "../rpcManager";

const responseSchema = Type.Object({
  title: Type.String(),
  score: Type.Number(),
});

describe("AIManager.generateObject", () => {
  it("returns schema-inferred data and serializes image inputs", async () => {
    let requestBody: AIGenerateRequest | undefined;
    const manager = createAIManager((body) => {
      requestBody = body;
      return { object: { title: "Diagram", score: 0.9 } };
    });

    const result: { title: string; score: number } =
      await manager.generateObject({
        prompt: "Describe these images",
        system: "Return a short label",
        model: "gpt-4.1-mini",
        schema: responseSchema,
        images: [
          {
            data: new Uint8Array([1, 2, 3]),
            mediaType: "image/png",
          },
          new Blob([new Uint8Array([4, 5])], { type: "image/webp" }),
          { url: new URL("https://example.com/image.jpg") },
        ],
      });

    expect(result).toEqual({ title: "Diagram", score: 0.9 });
    expect(requestBody).toEqual({
      prompt: "Describe these images",
      system: "Return a short label",
      model: "gpt-4.1-mini",
      schema: JSON.parse(JSON.stringify(responseSchema)),
      encodedSchema: encodeSchema(responseSchema),
      images: [
        { type: "data", data: "AQID", mediaType: "image/png" },
        { type: "data", data: "BAU=", mediaType: "image/webp" },
        { type: "url", url: "https://example.com/image.jpg" },
      ],
    });
  });

  it("rejects data that does not match the requested schema", async () => {
    const manager = createAIManager(() => ({
      object: { title: "Missing score" },
    }));

    await expect(
      manager.generateObject({
        prompt: "Describe the image",
        schema: responseSchema,
      })
    ).rejects.toThrow(
      "AI response does not match the requested schema: /score"
    );
  });

  it("validates prompts and image inputs before sending an RPC request", async () => {
    let requestCount = 0;
    const manager = createAIManager(() => {
      requestCount++;
      return { object: { title: "Unused", score: 0 } };
    });

    await expect(
      manager.generateObject({ prompt: " ", schema: responseSchema })
    ).rejects.toThrow("AI prompt must not be empty");
    await expect(
      manager.generateObject({
        prompt: "Describe this file",
        schema: responseSchema,
        images: [{ data: new Uint8Array([1]), mediaType: "application/pdf" }],
      })
    ).rejects.toThrow("Invalid AI image media type");
    await expect(
      manager.generateObject({
        prompt: "Describe this URL",
        schema: responseSchema,
        images: [{ url: "file:///tmp/image.png" }],
      })
    ).rejects.toThrow("AI image URLs must use http or https");

    expect(requestCount).toBe(0);
  });

  it("propagates RPC failures", async () => {
    const rpcManager = new RPCManager({ isConnected: true });
    const manager = new AIManager(rpcManager);

    rpcManager.addListener((request) => {
      rpcManager.handleMessage({
        id: request.id!,
        type: "error",
        error: "AI generation failed",
      });
    });

    await expect(
      manager.generateObject({
        prompt: "Describe the image",
        schema: responseSchema,
      })
    ).rejects.toBe("AI generation failed");
  });
});

describe("AIManager.generateText", () => {
  it("returns text and serializes images and model selection", async () => {
    let requestBody: AIGenerateTextRequest | undefined;
    const manager = createGenerationManager("/api/ai/generate/text", (body) => {
      requestBody = body as AIGenerateTextRequest;
      return { text: "A concise description" };
    });

    const result = await manager.generateText({
      prompt: "Describe this image",
      system: "Be concise",
      model: "gpt-4.1",
      images: [{ data: new Uint8Array([1, 2, 3]), mediaType: "image/png" }],
    });

    expect(result).toBe("A concise description");
    expect(requestBody).toEqual({
      prompt: "Describe this image",
      system: "Be concise",
      model: "gpt-4.1",
      images: [{ type: "data", data: "AQID", mediaType: "image/png" }],
    });
  });

  it("rejects an empty model", async () => {
    const manager = createGenerationManager("/api/ai/generate/text", () => ({
      text: "Unused",
    }));

    await expect(
      manager.generateText({ prompt: "Write text", model: " " })
    ).rejects.toThrow("AI model must not be empty");
  });
});

describe("AIManager.generateImage", () => {
  it("returns decoded image bytes and serializes model selection", async () => {
    let requestBody: AIGenerateImageRequest | undefined;
    const manager = createGenerationManager(
      "/api/ai/generate/image",
      (body) => {
        requestBody = body as AIGenerateImageRequest;
        return { base64: "AQID", mediaType: "image/png" };
      }
    );

    const result = await manager.generateImage({
      prompt: "Draw an otter",
      model: "gpt-image-2",
    });

    expect(result).toEqual({
      data: new Uint8Array([1, 2, 3]),
      mediaType: "image/png",
    });
    expect(requestBody).toEqual({
      prompt: "Draw an otter",
      model: "gpt-image-2",
    });
  });
});

function createAIManager(
  respond: (body: AIGenerateRequest) => { object: unknown }
) {
  const rpcManager = new RPCManager({ isConnected: true });
  const manager = new AIManager(rpcManager);

  rpcManager.addListener((request) => {
    if (
      request.url !== "/api/ai/generate" ||
      request.options.method !== "POST"
    ) {
      rpcManager.handleMessage({
        id: request.id!,
        type: "error",
        error: "Unexpected request",
      });
      return;
    }

    const body = JSON.parse(request.options.body!) as AIGenerateRequest;
    rpcManager.handleMessage({
      id: request.id!,
      type: "end",
      response: {
        status: 200,
        statusText: "OK",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(respond(body)),
      },
    });
  });

  return manager;
}

function createGenerationManager(
  url: "/api/ai/generate/text" | "/api/ai/generate/image",
  respond: (body: unknown) => unknown
) {
  const rpcManager = new RPCManager({ isConnected: true });
  const manager = new AIManager(rpcManager);

  rpcManager.addListener((request) => {
    if (request.url !== url || request.options.method !== "POST") {
      rpcManager.handleMessage({
        id: request.id!,
        type: "error",
        error: "Unexpected request",
      });
      return;
    }

    rpcManager.handleMessage({
      id: request.id!,
      type: "end",
      response: {
        status: 200,
        statusText: "OK",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(respond(JSON.parse(request.options.body ?? "{}"))),
      },
    });
  });

  return manager;
}