package prediction

import (
	"testing"

	"github.com/rose_cat707/luminary/internal/model"
)

func TestOpenAPISchemaForWorkflowTxt2Img(t *testing.T) {
	schema := OpenAPISchemaForWorkflow(model.WorkflowTxt2Img, "flux", "FLUX txt2img", nil)
	components, ok := schema["components"].(map[string]any)
	if !ok {
		t.Fatal("missing components")
	}
	schemas, ok := components["schemas"].(map[string]any)
	if !ok {
		t.Fatal("missing schemas")
	}
	input, ok := schemas["Input"].(map[string]any)
	if !ok {
		t.Fatal("missing Input schema")
	}
	props, ok := input["properties"].(map[string]any)
	if !ok {
		t.Fatal("missing properties")
	}
	if _, ok := props["prompt"]; !ok {
		t.Fatal("missing prompt property")
	}
	if _, ok := props["image"]; ok {
		t.Fatal("txt2img should not expose image input")
	}
	required, ok := input["required"].([]string)
	if !ok || len(required) == 0 || required[0] != "prompt" {
		t.Fatalf("unexpected required: %#v", input["required"])
	}
}

func TestOpenAPISchemaForWorkflowImg2Img(t *testing.T) {
	schema := OpenAPISchemaForWorkflow(model.WorkflowImg2Img, "flux-i2i", "FLUX img2img", nil)
	components := schema["components"].(map[string]any)
	schemas := components["schemas"].(map[string]any)
	input := schemas["Input"].(map[string]any)
	props := input["properties"].(map[string]any)
	if _, ok := props["image"]; !ok {
		t.Fatal("missing image property")
	}
	required := input["required"].([]string)
	if len(required) != 2 {
		t.Fatalf("expected prompt and image required, got %#v", required)
	}
}