iop/apps/edge/internal/openai/anthropic_bridge.go

package openai

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strings"

	"iop/packages/go/config"
)

type openAIChatBridgeResponse struct {
	ID      string `json:"id"`
	Model   string `json:"model"`
	Choices []struct {
		Message struct {
			Role             string `json:"role"`
			Content          any    `json:"content"`
			ReasoningContent string `json:"reasoning_content"`
			Reasoning        string `json:"reasoning"`
			ToolCalls        []struct {
				ID       string `json:"id"`
				Type     string `json:"type"`
				Function struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `json:"function"`
			} `json:"tool_calls"`
		} `json:"message"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
	Usage openAIChatBridgeUsage `json:"usage"`
}

type openAIChatBridgeUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	PromptDetails    struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`
}

type openAIChatBridgeError struct {
	Error struct {
		Type    string `json:"type"`
		Message string `json:"message"`
		Code    any    `json:"code"`
	} `json:"error"`
}

func prepareAnthropicChatBridge(body []byte, target string, profile config.ConcreteProtocolProfile) ([]byte, anthropicMessageRequest, error) {
	req, err := decodeAnthropicMessageRequest(body, true)
	if err != nil {
		return nil, req, err
	}
	if req.TopK != nil {
		return nil, req, fmt.Errorf("top_k is not supported by the Chat bridge")
	}
	if req.Thinking != nil && !profileSupportsAnthropicThinking(profile) {
		return nil, req, fmt.Errorf("selected Chat profile does not support thinking")
	}

	messages := make([]map[string]any, 0, len(req.Messages)+1)
	systemBlocks, err := decodeAnthropicSystem(req.System)
	if err != nil {
		return nil, req, err
	}
	if len(systemBlocks) > 0 {
		parts := make([]string, 0, len(systemBlocks))
		for _, block := range systemBlocks {
			parts = append(parts, block.Text)
		}
		messages = append(messages, map[string]any{"role": "system", "content": strings.Join(parts, "\n")})
	}

	for index, message := range req.Messages {
		blocks, err := decodeAnthropicContent(message.Content)
		if err != nil {
			return nil, req, fmt.Errorf("messages[%d].content: %w", index, err)
		}
		converted, err := anthropicMessageToChat(message.Role, blocks, profile)
		if err != nil {
			return nil, req, fmt.Errorf("messages[%d]: %w", index, err)
		}
		messages = append(messages, converted...)
	}

	chat := map[string]any{
		"model":      target,
		"messages":   messages,
		"max_tokens": *req.MaxTokens,
		"stream":     req.Stream,
	}
	if req.Stream {
		chat["stream_options"] = map[string]any{"include_usage": true}
	}
	if req.Temperature != nil {
		chat["temperature"] = *req.Temperature
	}
	if req.TopP != nil {
		chat["top_p"] = *req.TopP
	}
	if len(req.StopSequences) > 0 {
		chat["stop"] = req.StopSequences
	}
	if len(req.Metadata) > 0 && !bytes.Equal(bytes.TrimSpace(req.Metadata), []byte("null")) {
		var metadata map[string]any
		if err := json.Unmarshal(req.Metadata, &metadata); err != nil {
			return nil, req, fmt.Errorf("metadata must be an object")
		}
		chat["metadata"] = metadata
	}
	if len(req.Tools) > 0 {
		tools := make([]map[string]any, 0, len(req.Tools))
		for _, tool := range req.Tools {
			var schema map[string]any
			if err := json.Unmarshal(tool.InputSchema, &schema); err != nil {
				return nil, req, fmt.Errorf("tool %q input_schema is invalid", tool.Name)
			}
			function := map[string]any{"name": tool.Name, "parameters": schema}
			if tool.Description != "" {
				function["description"] = tool.Description
			}
			tools = append(tools, map[string]any{"type": "function", "function": function})
		}
		chat["tools"] = tools
	}
	if req.ToolChoice != nil {
		choice, parallel := anthropicToolChoiceToChat(*req.ToolChoice)
		chat["tool_choice"] = choice
		if parallel != nil {
			chat["parallel_tool_calls"] = *parallel
		}
	}
	if req.Thinking != nil {
		chat["think"] = true
		chat["include_reasoning"] = true
		chat["thinking_token_budget"] = req.Thinking.BudgetTokens
	}
	encoded, err := json.Marshal(chat)
	if err != nil {
		return nil, req, fmt.Errorf("encode Chat bridge request: %w", err)
	}
	return encoded, req, nil
}

func anthropicMessageToChat(role string, blocks []anthropicContentBlock, profile config.ConcreteProtocolProfile) ([]map[string]any, error) {
	if role == "assistant" {
		return anthropicAssistantToChat(blocks, profile)
	}
	var out []map[string]any
	var content []map[string]any
	flushContent := func() {
		if len(content) == 0 {
			return
		}
		out = append(out, map[string]any{"role": "user", "content": content})
		content = nil
	}
	for _, block := range blocks {
		switch block.Type {
		case "text":
			content = append(content, map[string]any{"type": "text", "text": block.Text})
		case "image":
			url := block.Source.URL
			if block.Source.Type == "base64" {
				url = "data:" + block.Source.MediaType + ";base64," + block.Source.Data
			}
			content = append(content, map[string]any{"type": "image_url", "image_url": map[string]any{"url": url}})
		case "tool_result":
			flushContent()
			text, err := anthropicToolResultText(block.Content)
			if err != nil {
				return nil, err
			}
			if block.IsError {
				text = "Error: " + text
			}
			out = append(out, map[string]any{"role": "tool", "tool_call_id": block.ToolUseID, "content": text})
		default:
			return nil, fmt.Errorf("content block %q is invalid for a user message", block.Type)
		}
	}
	flushContent()
	if len(out) == 0 {
		return nil, fmt.Errorf("user message content is empty")
	}
	return out, nil
}

func anthropicAssistantToChat(blocks []anthropicContentBlock, profile config.ConcreteProtocolProfile) ([]map[string]any, error) {
	message := map[string]any{"role": "assistant"}
	var content []map[string]any
	var reasoning []string
	var toolCalls []map[string]any
	for _, block := range blocks {
		switch block.Type {
		case "text":
			content = append(content, map[string]any{"type": "text", "text": block.Text})
		case "thinking":
			if !profileSupportsAnthropicThinking(profile) {
				return nil, fmt.Errorf("selected Chat profile does not support thinking blocks")
			}
			if block.Signature != "" {
				return nil, fmt.Errorf("signed thinking blocks cannot be represented by the Chat bridge")
			}
			reasoning = append(reasoning, block.Thinking)
		case "tool_use":
			toolCalls = append(toolCalls, map[string]any{
				"id": block.ID, "type": "function",
				"function": map[string]any{"name": block.Name, "arguments": string(block.Input)},
			})
		default:
			return nil, fmt.Errorf("content block %q is invalid for an assistant message", block.Type)
		}
	}
	if len(content) > 0 {
		message["content"] = content
	} else {
		message["content"] = nil
	}
	if len(reasoning) > 0 {
		message["reasoning_content"] = strings.Join(reasoning, "")
	}
	if len(toolCalls) > 0 {
		message["tool_calls"] = toolCalls
	}
	return []map[string]any{message}, nil
}

func anthropicToolResultText(raw json.RawMessage) (string, error) {
	if len(bytes.TrimSpace(raw)) == 0 || bytes.Equal(bytes.TrimSpace(raw), []byte("null")) {
		return "", nil
	}
	var text string
	if err := json.Unmarshal(raw, &text); err == nil {
		return text, nil
	}
	blocks, err := decodeAnthropicContent(raw)
	if err != nil {
		return "", err
	}
	parts := make([]string, 0, len(blocks))
	for _, block := range blocks {
		if block.Type != "text" {
			return "", fmt.Errorf("Chat bridge tool_result supports only text content")
		}
		parts = append(parts, block.Text)
	}
	return strings.Join(parts, ""), nil
}

func anthropicToolChoiceToChat(choice anthropicToolChoice) (any, *bool) {
	parallel := !choice.DisableParallelToolUse
	switch choice.Type {
	case "auto":
		return "auto", &parallel
	case "any":
		return "required", &parallel
	case "none":
		return "none", &parallel
	default:
		return map[string]any{"type": "function", "function": map[string]any{"name": choice.Name}}, &parallel
	}
}

func profileSupportsAnthropicThinking(profile config.ConcreteProtocolProfile) bool {
	for _, key := range []string{"thinking", "reasoning"} {
		value, ok := profile.Extensions[key]
		if enabled, valid := value.(bool); ok && valid && enabled {
			return true
		}
	}
	return false
}

func convertChatResponseToAnthropic(body []byte, requestModel string) (anthropicMessageResponse, error) {
	var response openAIChatBridgeResponse
	if err := json.Unmarshal(body, &response); err != nil {
		return anthropicMessageResponse{}, fmt.Errorf("decode Chat response: %w", err)
	}
	if len(response.Choices) != 1 {
		return anthropicMessageResponse{}, fmt.Errorf("Chat response must contain exactly one choice")
	}
	choice := response.Choices[0]
	content := make([]map[string]any, 0, 2+len(choice.Message.ToolCalls))
	reasoning := choice.Message.ReasoningContent
	if reasoning == "" {
		reasoning = choice.Message.Reasoning
	}
	if reasoning != "" {
		content = append(content, map[string]any{"type": "thinking", "thinking": reasoning, "signature": ""})
	}
	text, err := openAIChatContentText(choice.Message.Content)
	if err != nil {
		return anthropicMessageResponse{}, err
	}
	if text != "" {
		content = append(content, map[string]any{"type": "text", "text": text})
	}
	for _, call := range choice.Message.ToolCalls {
		if call.ID == "" || call.Function.Name == "" || !json.Valid([]byte(call.Function.Arguments)) {
			return anthropicMessageResponse{}, fmt.Errorf("Chat tool call has invalid id, name, or arguments")
		}
		var input any
		if err := json.Unmarshal([]byte(call.Function.Arguments), &input); err != nil {
			return anthropicMessageResponse{}, fmt.Errorf("decode Chat tool arguments: %w", err)
		}
		content = append(content, map[string]any{"type": "tool_use", "id": call.ID, "name": call.Function.Name, "input": input})
	}
	stopReason, err := anthropicStopReason(choice.FinishReason)
	if err != nil {
		return anthropicMessageResponse{}, err
	}
	id := response.ID
	if id == "" {
		id = "msg_iop"
	}
	return anthropicMessageResponse{
		ID: id, Type: "message", Role: "assistant", Model: requestModel,
		Content: content, StopReason: stopReason, StopSequence: nil,
		Usage: anthropicUsage{
			InputTokens: response.Usage.PromptTokens, OutputTokens: response.Usage.CompletionTokens,
			CacheReadInputTokens: response.Usage.PromptDetails.CachedTokens,
		},
	}, nil
}

func openAIChatContentText(content any) (string, error) {
	switch value := content.(type) {
	case nil:
		return "", nil
	case string:
		return value, nil
	case []any:
		var parts []string
		for _, item := range value {
			block, ok := item.(map[string]any)
			if !ok || block["type"] != "text" {
				return "", fmt.Errorf("unsupported Chat response content block")
			}
			text, ok := block["text"].(string)
			if !ok {
				return "", fmt.Errorf("Chat response text block is invalid")
			}
			parts = append(parts, text)
		}
		return strings.Join(parts, ""), nil
	default:
		return "", fmt.Errorf("unsupported Chat response content")
	}
}

func anthropicStopReason(finishReason *string) (*string, error) {
	if finishReason == nil || *finishReason == "" {
		return nil, nil
	}
	var reason string
	switch *finishReason {
	case "stop":
		reason = "end_turn"
	case "length":
		reason = "max_tokens"
	case "tool_calls":
		reason = "tool_use"
	default:
		return nil, fmt.Errorf("unsupported Chat finish_reason %q", *finishReason)
	}
	return &reason, nil
}

func convertChatErrorToAnthropic(body []byte) anthropicErrorResponse {
	var provider openAIChatBridgeError
	if json.Unmarshal(body, &provider) == nil && strings.TrimSpace(provider.Error.Message) != "" {
		errorType := strings.TrimSpace(provider.Error.Type)
		if errorType == "" {
			errorType = "api_error"
		}
		return anthropicErrorResponse{Type: "error", Error: errorBody{Type: errorType, Message: provider.Error.Message}}
	}
	return anthropicErrorResponse{Type: "error", Error: errorBody{Type: "api_error", Message: "upstream provider error"}}
}