package vision

import (
	"context"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

func TestMessageResolver_FoldsImageToDescription(t *testing.T) {
	t.Parallel()
	rt := NewRouter(WithVisionTool("vision"), WithOCRTool("ocr"))
	invoked := ""
	res := NewMessageResolver(rt, func() ModelInfo {
		return ModelInfo{Name: "text-model", SupportsVision: false}
	}, func(_ context.Context, name string, args map[string]any) (any, error) {
		invoked = name
		return map[string]any{"description": "一只橘色的猫"}, nil
	})

	msg, err := res.Resolve(context.Background(), model.ImageMessage(model.RoleUser, "看到了吗", "data:image/png;base64,AAAA"))
	if err != nil {
		t.Fatalf("Resolve: %v", err)
	}
	if invoked != "vision" {
		t.Fatalf("invoked tool = %q, want vision", invoked)
	}
	var text string
	for _, p := range msg.Content {
		if p.Type == model.ContentTypeText {
			text += p.Text
		}
		if p.Type == model.ContentTypeImage {
			t.Fatalf("image part still present after tool route: %+v", p)
		}
	}
	if text == "" || text == "看到了吗" {
		t.Fatalf("expected a description fragment, got %q", text)
	}
}

func TestMessageResolver_KeepsImageForVisionModel(t *testing.T) {
	t.Parallel()
	rt := NewRouter(WithVisionTool("vision"), WithOCRTool("ocr"))
	res := NewMessageResolver(rt, func() ModelInfo {
		return ModelInfo{Name: "vision-model", SupportsVision: true}
	}, func(_ context.Context, _ string, _ map[string]any) (any, error) {
		t.Fatalf("tool should not be invoked when the main model is sighted")
		return nil, nil
	})

	msg, err := res.Resolve(context.Background(), model.ImageMessage(model.RoleUser, "x", "data:image/png;base64,AAAA"))
	if err != nil {
		t.Fatalf("Resolve: %v", err)
	}
	if !messageHasImagePart(msg) {
		t.Fatalf("expected the image part to be kept for a vision-capable main model")
	}
}

func TestMessageResolver_OCRToolForText(t *testing.T) {
	t.Parallel()
	rt := NewRouter(WithVisionTool("vision"), WithOCRTool("ocr"))
	invoked := ""
	res := NewMessageResolver(rt, func() ModelInfo {
		return ModelInfo{Name: "text-model", SupportsVision: false}
	}, func(_ context.Context, name string, _ map[string]any) (any, error) {
		invoked = name
		return map[string]any{"text": "HELLO"}, nil
	})

	// 通过 router 的路由请求路径强制给出文本优先提示：这里直接经 Resolve
	// 走通，而 Resolve 不传任何提示，因此优先使用视觉工具；
	// 断言至少有一个文本部件替换了图像。
	msg, err := res.Resolve(context.Background(), model.ImageMessage(model.RoleUser, "x", "data:image/png;base64,AAAA"))
	if err != nil {
		t.Fatalf("Resolve: %v", err)
	}
	if invoked != "vision" {
		t.Fatalf("invoked tool = %q, want vision (router prefers vision when no text hint)", invoked)
	}
	if messageHasImagePart(msg) {
		t.Fatalf("expected the image to be replaced by text")
	}
}

func messageHasImagePart(m model.Message) bool {
	for _, p := range m.Content {
		if p.Type == model.ContentTypeImage {
			return true
		}
	}
	return false
}