import { expect, test } from "@playwright/test";
import {
  createAtlasCloudImagePreset,
  inferAtlasCloudMediaKind,
} from "@/lib/ai/atlascloud-models";
import {
  is3DModelDescriptor,
  isImageGenerationModelDescriptor,
  isVideoModelDescriptor,
} from "@/lib/ai/model-classification";

const atlasHunyuan3D = {
  model: "tencent/hunyuan3d-pro/text-to-3d",
  type: "Image",
  displayName: "Hunyuan 3D Pro Text-to-3D",
  categories: ["TEXT-TO-3D"],
  tags: ["TEXT-TO-3D", "NEW"],
};

test.describe("AI model media classification", () => {
  test("keeps Atlas 3D entries out of image and video categories", () => {
    expect(inferAtlasCloudMediaKind(atlasHunyuan3D)).toBe("3d");

    const descriptor = {
      providerType: "atlascloud",
      modelId: atlasHunyuan3D.model,
      capabilities: {
        kind: "3d",
        mediaType: "3d",
        supportedModes: ["text-to-3d"],
      },
      supportsImageGen: false,
    };

    expect(is3DModelDescriptor(descriptor)).toBe(true);
    expect(isImageGenerationModelDescriptor(descriptor)).toBe(false);
    expect(isVideoModelDescriptor(descriptor)).toBe(false);
  });

  test("writes 3D Atlas presets with model-level capabilities", () => {
    const preset = createAtlasCloudImagePreset(
      atlasHunyuan3D,
      {
        components: {
          schemas: {
            Input: {
              type: "object",
              properties: {
                model: { type: "string" },
                prompt: { type: "string" },
              },
              required: ["model", "prompt"],
            },
          },
        },
      },
      0
    );

    expect(preset.supportsImageGen).toBe(false);
    expect(preset.capabilities).toMatchObject({
      kind: "3d",
      mediaType: "3d",
      supportedModes: ["text-to-3d"],
    });
  });

  test("still distinguishes Atlas image and video descriptors", () => {
    const imageDescriptor = {
      providerType: "atlascloud",
      modelId: "black-forest-labs/flux-2-pro/text-to-image",
      capabilities: {
        kind: "text-to-image",
        mediaType: "image",
        supportedModes: ["text-to-image"],
      },
      supportsImageGen: true,
    };
    const videoDescriptor = {
      providerType: "atlascloud",
      modelId: "veed/fabric-1.0/image-to-video",
      capabilities: {
        kind: "video",
        mediaType: "video",
        supportedModes: ["image-to-video"],
      },
      supportsImageGen: false,
    };

    expect(isImageGenerationModelDescriptor(imageDescriptor)).toBe(true);
    expect(isVideoModelDescriptor(imageDescriptor)).toBe(false);
    expect(isVideoModelDescriptor(videoDescriptor)).toBe(true);
    expect(isImageGenerationModelDescriptor(videoDescriptor)).toBe(false);
  });
});