package vision

import (
	"context"
	"fmt"

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

// ModelInfoProvider supplies the main model's vision capability so the resolver
// can route an image the way the vision Router decides (design/06 §0). It is
// called per image; the same provider is shared across a single resolution.
type ModelInfoProvider func() ModelInfo

// ToolInvoker invokes a registered tool (the vision or OCR tool) by name with
// the given arguments. It is satisfied by *tool.Dispatcher.Invoke.
type ToolInvoker func(ctx context.Context, name string, args map[string]any) (any, error)

// MessageResolver is the loop's concrete ImageResolver. For each image part in a
// user message it asks the vision Router how to route it: when the main model
// supports vision and is confident, the image part is kept so the main model sees
// the image; otherwise the vision/OCR tool is invoked and its text result replaces
// the image part, so a text-only main model can reason over the image ("先看图再
// 对话"). It is safe for concurrent use.
type MessageResolver struct {
	router *Router
	info   ModelInfoProvider
	invoke ToolInvoker
}

// NewMessageResolver builds a MessageResolver. router decides the layer; info
// reports the main model's vision capability; invoke runs a registered vision/OCR
// tool. A nil invoke leaves any tool route failed (the image is then kept), so a
// resolver without tools degrades to the main-model path.
func NewMessageResolver(router *Router, info ModelInfoProvider, invoke ToolInvoker) *MessageResolver {
	return &MessageResolver{
		router: router,
		info:   info,
		invoke: invoke,
	}
}

// Resolve rewrites msg so every image part is either kept (main model sees it) or
// replaced by a text description produced by the vision/OCR tool. Text parts are
// preserved in order. It errors only when an image cannot be routed or a tool
// invocation fails; the caller (loop) then keeps the original message.
func (r *MessageResolver) Resolve(ctx context.Context, msg model.Message) (model.Message, error) {
	if r.router == nil {
		return msg, nil
	}
	out := make([]model.Content, 0, len(msg.Content))
	for _, part := range msg.Content {
		if part.Type != model.ContentTypeImage {
			out = append(out, part)
			continue
		}
		replacement, err := r.resolveImage(ctx, part.ImageURL)
		if err != nil {
			return model.Message{}, err
		}
		out = append(out, replacement)
	}
	return model.Message{Role: msg.Role, Content: out}, nil
}

// resolveImage picks a route and either keeps the image (KindMain) or folds it
// into a text description via the vision/OCR tool.
func (r *MessageResolver) resolveImage(ctx context.Context, imageURL string) (model.Content, error) {
	target := ModelInfo{}
	if r.info != nil {
		target = r.info()
	}
	decision := r.router.Route(target, RouteRequest{})
	switch decision.Kind {
	case KindMain:
		// A vision-capable main model sees the image directly.
		return model.Content{Type: model.ContentTypeImage, ImageURL: imageURL}, nil
	case KindVisionTool:
		return r.invokeTool(ctx, decision.Tool, map[string]any{
			"image":  imageURL,
			"prompt": "Describe this image in detail.",
		}, "description")
	case KindOCRTool:
		return r.invokeTool(ctx, decision.Tool, map[string]any{
			"image": imageURL,
		}, "text")
	default:
		return model.Content{}, fmt.Errorf("image route: %w: %s", ErrUnsupportedSource, decision.Reason)
	}
}

// invokeTool runs a registered vision/OCR tool and turns its text result into a
// text content part. field is the result map key ("description" or "text").
func (r *MessageResolver) invokeTool(ctx context.Context, toolName string, args map[string]any, field string) (model.Content, error) {
	if r.invoke == nil {
		return model.Content{}, fmt.Errorf("%s: %w", toolName, ErrNotConfigured)
	}
	result, err := r.invoke(ctx, toolName, args)
	if err != nil {
		return model.Content{}, fmt.Errorf("image tool %s: %w", toolName, err)
	}
	return model.Content{Type: model.ContentTypeText, Text: descriptionFragment(fieldText(result, field))}, nil
}

// fieldText extracts a string field from a tool result map, tolerating the
// map[string]any shape returned by the vision/OCR tools.
func fieldText(result any, field string) string {
	m, ok := result.(map[string]any)
	if !ok {
		return ""
	}
	s, _ := m[field].(string)
	return s
}

// descriptionFragment labels a tool-produced description so the model can tell it
// refers to the picture the user attached.
func descriptionFragment(desc string) string {
	if desc == "" {
		return "[图片内容] 无法识别该图片"
	}
	return "[图片内容] " + desc
}