// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// anthropic-proxy exposes an Anthropic Messages API-compatible HTTP server
// that forwards requests to the CMaaS confidential inference service via
// the CMaaS SDK.
//
// Usage:
//
//	export DASHSCOPE_API_KEY=<your-api-key>
//	./anthropic-proxy \
//	    --endpoint wss://<modelstudio-endpoint>/api-ws/v1/inference \
//	    --model <model-name> \
//	    --listen :8081 \
//	    --workers 16
//
// Then point any Anthropic-SDK client at http://localhost:8081:
//
//	client = anthropic.Anthropic(base_url="http://localhost:8081", api_key="ignored")
//	client.messages.create(model="...", max_tokens=64,
//	    messages=[{"role":"user","content":"hello"}])
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/dashscope/dashscope-confidential-maas/sdk"

	"github.com/google/uuid"
)

// ---- Anthropic Messages API request / response types ----

type anthropicRequest struct {
	Model         string             `json:"model"`
	Messages      []anthropicMessage `json:"messages"`
	MaxTokens     int                `json:"max_tokens"`
	System        json.RawMessage    `json:"system,omitempty"`
	Temperature   *float64           `json:"temperature,omitempty"`
	TopP          *float64           `json:"top_p,omitempty"`
	TopK          *int               `json:"top_k,omitempty"`
	StopSequences []string           `json:"stop_sequences,omitempty"`
	Stream        bool               `json:"stream,omitempty"`
	Thinking      *anthropicThinking `json:"thinking,omitempty"`
	Tools         json.RawMessage    `json:"tools,omitempty"`
	ToolChoice    json.RawMessage    `json:"tool_choice,omitempty"`
}

type anthropicThinking struct {
	Type         string `json:"type"`
	BudgetTokens *int   `json:"budget_tokens,omitempty"`
}

type anthropicMessage struct {
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
}

type anthropicTextBlock struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

// anthropicToolUseBlock is emitted in the response content[] when the model
// invoked a tool. Input is a JSON object reconstructed from the upstream
// OpenAI-style tool_calls[].function.arguments (which is a JSON string).
type anthropicToolUseBlock struct {
	Type  string         `json:"type"`
	ID    string         `json:"id"`
	Name  string         `json:"name"`
	Input map[string]any `json:"input"`
}

type anthropicUsage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

// anthropicResponse mirrors the openai-proxy convention of stashing
// reasoning in a non-spec top-level side field (ReasoningContent) rather
// than as an unsigned thinking block inside Content[]. The official
// anthropic SDK refuses to parse thinking blocks without a `signature`,
// which our upstream does not produce; the SDK's Message model accepts
// unknown top-level fields (extra="allow"), so the side field is safe.
type anthropicResponse struct {
	ID               string         `json:"id"`
	Type             string         `json:"type"`
	Role             string         `json:"role"`
	Content          []any          `json:"content"` // mix of anthropicTextBlock and anthropicToolUseBlock
	Model            string         `json:"model"`
	StopReason       *string        `json:"stop_reason"`
	StopSequence     *string        `json:"stop_sequence"`
	Usage            anthropicUsage `json:"usage"`
	ReasoningContent string         `json:"reasoning_content,omitempty"`
	RequestID        string         `json:"request_id,omitempty"`
}

// ---- Field validation ----

var supportedMessagesRequestFields = map[string]struct{}{
	"model":          {},
	"messages":       {},
	"max_tokens":     {},
	"system":         {},
	"temperature":    {},
	"top_p":          {},
	"top_k":          {},
	"stop_sequences": {},
	"stream":         {},
	"thinking":       {},
	"tools":          {},
	"tool_choice":    {},
}

var unsupportedSemanticMessagesRequestFields = map[string]struct{}{
	"container":    {},
	"mcp_servers":  {},
	"top_logprobs": {},
}

var compatiblePlatformMessagesRequestFields = map[string]struct{}{
	"metadata":     {},
	"service_tier": {},
	// Claude Code CLI emits these alongside the standard Anthropic fields:
	// `output_config` is a structured-output hint, `context_management` is
	// a context-window pruning policy. The cmaas proxy doesn't drive either
	// at this layer, so they're no-ops here.
	"output_config":      {},
	"context_management": {},
}

type requestFieldError struct {
	Field   string
	Message string
}

func (e *requestFieldError) Error() string { return e.Message }

func validateMessagesRequestFields(body []byte) error {
	var fields map[string]json.RawMessage
	if err := json.Unmarshal(body, &fields); err != nil {
		return err
	}
	for field := range fields {
		if _, ok := supportedMessagesRequestFields[field]; ok {
			continue
		}
		if _, ok := compatiblePlatformMessagesRequestFields[field]; ok {
			log.Printf("[INFO] ignoring Anthropic platform field field=%s", field)
			continue
		}
		if _, ok := unsupportedSemanticMessagesRequestFields[field]; ok {
			return &requestFieldError{Field: field, Message: fmt.Sprintf("unsupported messages field %q", field)}
		}
		return &requestFieldError{Field: field, Message: fmt.Sprintf("unknown messages field %q", field)}
	}
	return nil
}

// ---- Client pool ----

type pool struct {
	ch chan *sdk.Client
}

func newPool(n int, opts *sdk.Options) (*pool, error) {
	p := &pool{ch: make(chan *sdk.Client, n)}
	for i := 0; i < n; i++ {
		c, err := sdk.NewClient(opts)
		if err != nil {
			return nil, fmt.Errorf("failed to create client %d: %w", i, err)
		}
		p.ch <- c
	}
	return p, nil
}

func (p *pool) acquire() *sdk.Client  { return <-p.ch }
func (p *pool) release(c *sdk.Client) { p.ch <- c }

func (p *pool) close() {
	close(p.ch)
	for c := range p.ch {
		c.Close()
	}
}

// ---- HTTP handler ----

type handler struct {
	pool  *pool
	model string
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.URL.Path {
	case "/v1/messages", "/messages":
		h.messages(w, r)
	case "/v1/models", "/models":
		h.listModels(w, r)
	default:
		http.NotFound(w, r)
	}
}

func (h *handler) listModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeAnthropicError(w, http.StatusMethodNotAllowed, "invalid_request_error", "method not allowed", "")
		return
	}
	resp := map[string]interface{}{
		"data": []map[string]interface{}{
			{
				"type":         "model",
				"id":           h.model,
				"display_name": h.model,
				"created_at":   time.Now().UTC().Format(time.RFC3339),
			},
		},
		"has_more": false,
		"first_id": h.model,
		"last_id":  h.model,
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

func (h *handler) messages(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeAnthropicError(w, http.StatusMethodNotAllowed, "invalid_request_error", "method not allowed", "")
		return
	}

	body, err := io.ReadAll(r.Body)
	if err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "failed to read request body", "")
		return
	}
	defer r.Body.Close()

	if err := validateMessagesRequestFields(body); err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error(), "")
		return
	}

	var req anthropicRequest
	if err := json.Unmarshal(body, &req); err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "invalid JSON: "+err.Error(), "")
		return
	}

	if req.MaxTokens <= 0 {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "max_tokens is required and must be positive", "")
		return
	}

	msgs, err := buildSDKMessages(req.System, req.Messages)
	if err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error(), "")
		return
	}

	requestID := uuid.New().String()
	respID := "msg_" + requestID
	sdkReq := &sdk.ChatCompletionRequest{
		RequestID: requestID,
		Messages:  msgs,
		Stream:    req.Stream,
		MaxTokens: sdk.Int(req.MaxTokens),
	}
	if err := applyMessageParameters(&req, sdkReq); err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error(), requestID)
		return
	}

	log.Printf("[INFO] request started id=%s endpoint=/v1/messages stream=%t backend_model=%q", respID, req.Stream, h.model)

	client := h.pool.acquire()
	defer h.pool.release(client)

	ctx := r.Context()
	completion, err := client.Chat(ctx, sdkReq)
	if err != nil {
		log.Printf("[WARN] request failed id=%s task_id=%s error=%v", respID, requestID, err)
		writeAnthropicError(w, http.StatusInternalServerError, "api_error", "inference error: "+err.Error(), requestID)
		return
	}
	log.Printf("[INFO] cmaas request mapped id=%s task_id=%s stream=%t", respID, completion.TaskID, req.Stream)

	// When the client sent tool_choice="none", defend against a backend
	// that nonetheless returns tool calls — strip them at the proxy.
	toolChoiceNone := isToolChoiceNone(sdkReq.ToolChoice)

	if req.Stream {
		h.writeStream(w, completion, respID, requestID, h.model, toolChoiceNone)
	} else {
		h.writeResponse(w, completion, respID, requestID, h.model, toolChoiceNone)
	}
}

// isToolChoiceNone returns true when the serialized tool_choice value is
// the JSON string "none" (the canonical OpenAI shape after Anthropic's
// translation — `{type:"none"}` and `"none"` both collapse to `"none"`).
func isToolChoiceNone(raw json.RawMessage) bool {
	if len(raw) == 0 {
		return false
	}
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return s == "none"
	}
	return false
}

// ensureToolCallIDs returns a copy of the slice with any empty ID filled
// in. Anthropic's tool_use.id is required for the client to correlate the
// subsequent tool_result; missing ids would surface as schema violations
// in the official SDK.
func ensureToolCallIDs(calls []sdk.ToolCall) []sdk.ToolCall {
	if len(calls) == 0 {
		return nil
	}
	out := make([]sdk.ToolCall, len(calls))
	copy(out, calls)
	for i := range out {
		if out[i].ID == "" {
			out[i].ID = "toolu_" + uuid.New().String()
		}
	}
	return out
}

func applyMessageParameters(req *anthropicRequest, sdkReq *sdk.ChatCompletionRequest) error {
	sdkReq.Temperature = req.Temperature
	sdkReq.TopP = req.TopP
	sdkReq.TopK = req.TopK
	if len(req.StopSequences) > 0 {
		sdkReq.Stop = req.StopSequences
	}

	if hasJSONValue(req.Tools) {
		tools, err := translateAnthropicTools(req.Tools)
		if err != nil {
			return err
		}
		sdkReq.Tools = tools
	}
	if hasJSONValue(req.ToolChoice) {
		choice, err := translateAnthropicToolChoice(req.ToolChoice)
		if err != nil {
			return err
		}
		sdkReq.ToolChoice = choice
	}

	if req.Thinking == nil {
		return nil
	}
	switch req.Thinking.Type {
	case "enabled", "adaptive":
		// `adaptive` is Claude Code's variable-budget mode. The upstream
		// chat template has no equivalent, so leave the budget to the model.
		sdkReq.ChatTemplateKwargs = map[string]interface{}{"enable_thinking": true}
		if req.Thinking.BudgetTokens != nil {
			log.Printf("[INFO] thinking.budget_tokens=%d is not propagated to the upstream chat template", *req.Thinking.BudgetTokens)
		}
	case "disabled":
		sdkReq.ChatTemplateKwargs = map[string]interface{}{"enable_thinking": false}
	default:
		return fmt.Errorf("thinking.type must be one of \"enabled\", \"adaptive\", \"disabled\" (got %q)", req.Thinking.Type)
	}
	return nil
}

// translateAnthropicTools converts Anthropic `tools` array into the
// OpenAI-style sdk.Tool slice expected by the CMaaS SDK.
//
//	Anthropic: {name, description, input_schema}
//	OpenAI:    {type:"function", function:{name, description, parameters}}
func translateAnthropicTools(raw json.RawMessage) ([]sdk.Tool, error) {
	var in []struct {
		Name        string                 `json:"name"`
		Description string                 `json:"description"`
		InputSchema map[string]interface{} `json:"input_schema"`
	}
	if err := json.Unmarshal(raw, &in); err != nil {
		return nil, fmt.Errorf("tools must be an array of {name, description, input_schema} objects: %w", err)
	}
	out := make([]sdk.Tool, len(in))
	for i, t := range in {
		if t.Name == "" {
			return nil, fmt.Errorf("tools[%d] missing name", i)
		}
		out[i] = sdk.Tool{
			Type: "function",
			Function: sdk.FunctionSchema{
				Name:        t.Name,
				Description: t.Description,
				Parameters:  t.InputSchema,
			},
		}
	}
	return out, nil
}

// translateAnthropicToolChoice converts Anthropic `tool_choice` to OpenAI's
// form. Anthropic accepts:
//   - "auto" / "any" / "none"  → OpenAI "auto" / "required" / "none"
//   - {"type":"tool", "name":"X"} → {"type":"function","function":{"name":"X"}}
//   - {"type":"auto"|"any"|"none"} (object form) → corresponding string
//
// `disable_parallel_tool_use` is dropped (no per-call equivalent in OpenAI).
func translateAnthropicToolChoice(raw json.RawMessage) (json.RawMessage, error) {
	var asString string
	if err := json.Unmarshal(raw, &asString); err == nil {
		switch asString {
		case "auto", "none":
			return json.RawMessage(fmt.Sprintf("%q", asString)), nil
		case "any":
			return json.RawMessage(`"required"`), nil
		default:
			return nil, fmt.Errorf("tool_choice string must be one of: auto, any, none")
		}
	}
	var obj struct {
		Type string `json:"type"`
		Name string `json:"name,omitempty"`
	}
	if err := json.Unmarshal(raw, &obj); err != nil {
		return nil, fmt.Errorf("tool_choice must be a string or object: %w", err)
	}
	switch obj.Type {
	case "auto", "none":
		return json.RawMessage(fmt.Sprintf("%q", obj.Type)), nil
	case "any":
		return json.RawMessage(`"required"`), nil
	case "tool":
		if obj.Name == "" {
			return nil, fmt.Errorf("tool_choice type=tool requires name")
		}
		out, _ := json.Marshal(map[string]any{
			"type":     "function",
			"function": map[string]any{"name": obj.Name},
		})
		return out, nil
	default:
		return nil, fmt.Errorf("tool_choice type %q is not supported", obj.Type)
	}
}

// hasJSONValue returns true when raw is non-empty and not the JSON literal null.
func hasJSONValue(raw json.RawMessage) bool {
	t := strings.TrimSpace(string(raw))
	return t != "" && t != "null"
}

// parsedContentBlock is a typed view of an Anthropic content block. The
// proxy supports text on every message role, tool_use only on assistant,
// and tool_result only on user. Other block types (image, etc.) are
// rejected upstream of this struct.
type parsedContentBlock struct {
	BlockType string

	Text string // for text blocks

	// tool_use (assistant emitting a function call as input context)
	ToolID    string
	ToolName  string
	ToolInput json.RawMessage // raw JSON object; passed through to OpenAI as Arguments string

	// tool_result (user feeding back a tool execution result)
	ToolUseID  string
	ResultText string // concatenated text from content (string or text-block array)
}

// extractText reads system-style content (string or array of text blocks)
// and returns concatenated text. Only text blocks are supported; tool_use
// or tool_result inside `system` would be a spec violation and is rejected.
func extractText(raw json.RawMessage, where string) (string, error) {
	var asString string
	if err := json.Unmarshal(raw, &asString); err == nil {
		return asString, nil
	}
	blocks, err := parseContentBlocks(raw, where)
	if err != nil {
		return "", err
	}
	var sb strings.Builder
	for j, b := range blocks {
		if b.BlockType != "text" {
			return "", fmt.Errorf("%s[%d] type %q is not supported here (only text blocks are supported)", where, j, b.BlockType)
		}
		sb.WriteString(b.Text)
	}
	return sb.String(), nil
}

// parseContentBlocks accepts an Anthropic content field that is either a
// JSON string or a JSON array of content blocks. String form is normalized
// into a single text block. The image block type is rejected; tool_use and
// tool_result are returned typed (role validation happens at caller).
func parseContentBlocks(raw json.RawMessage, where string) ([]parsedContentBlock, error) {
	var asString string
	if err := json.Unmarshal(raw, &asString); err == nil {
		return []parsedContentBlock{{BlockType: "text", Text: asString}}, nil
	}

	var rawBlocks []map[string]json.RawMessage
	if err := json.Unmarshal(raw, &rawBlocks); err != nil {
		return nil, fmt.Errorf("%s must be a string or an array of content blocks", where)
	}

	out := make([]parsedContentBlock, 0, len(rawBlocks))
	for j, b := range rawBlocks {
		var btype string
		if err := json.Unmarshal(b["type"], &btype); err != nil {
			return nil, fmt.Errorf("%s[%d] missing block type", where, j)
		}
		switch btype {
		case "text":
			var text string
			if err := json.Unmarshal(b["text"], &text); err != nil {
				return nil, fmt.Errorf("%s[%d] missing text field", where, j)
			}
			out = append(out, parsedContentBlock{BlockType: "text", Text: text})
		case "tool_use":
			pb := parsedContentBlock{BlockType: "tool_use"}
			if err := json.Unmarshal(b["id"], &pb.ToolID); err != nil {
				return nil, fmt.Errorf("%s[%d] tool_use missing id", where, j)
			}
			if err := json.Unmarshal(b["name"], &pb.ToolName); err != nil {
				return nil, fmt.Errorf("%s[%d] tool_use missing name", where, j)
			}
			if input, ok := b["input"]; ok && len(input) > 0 && string(input) != "null" {
				pb.ToolInput = input
			} else {
				pb.ToolInput = json.RawMessage("{}")
			}
			out = append(out, pb)
		case "tool_result":
			pb := parsedContentBlock{BlockType: "tool_result"}
			if err := json.Unmarshal(b["tool_use_id"], &pb.ToolUseID); err != nil {
				return nil, fmt.Errorf("%s[%d] tool_result missing tool_use_id", where, j)
			}
			text, err := extractToolResultContent(b["content"], fmt.Sprintf("%s[%d].content", where, j))
			if err != nil {
				return nil, err
			}
			pb.ResultText = text
			out = append(out, pb)
		default:
			return nil, fmt.Errorf("%s[%d] type %q is not supported", where, j, btype)
		}
	}
	return out, nil
}

// extractToolResultContent normalizes the tool_result.content field, which
// per the Anthropic spec may be a plain string or an array containing text
// blocks. Image-typed blocks inside tool_result are rejected (no vision on
// this CMaaS path).
func extractToolResultContent(raw json.RawMessage, where string) (string, error) {
	if len(raw) == 0 || string(raw) == "null" {
		return "", nil
	}
	var asString string
	if err := json.Unmarshal(raw, &asString); err == nil {
		return asString, nil
	}
	var rawBlocks []map[string]json.RawMessage
	if err := json.Unmarshal(raw, &rawBlocks); err != nil {
		return "", fmt.Errorf("%s must be a string or an array of content blocks", where)
	}
	var sb strings.Builder
	for j, b := range rawBlocks {
		var btype string
		if err := json.Unmarshal(b["type"], &btype); err != nil {
			return "", fmt.Errorf("%s[%d] missing block type", where, j)
		}
		if btype != "text" {
			return "", fmt.Errorf("%s[%d] type %q is not supported in tool_result.content", where, j, btype)
		}
		var text string
		if err := json.Unmarshal(b["text"], &text); err != nil {
			return "", fmt.Errorf("%s[%d] missing text field", where, j)
		}
		sb.WriteString(text)
	}
	return sb.String(), nil
}

// buildSDKMessages converts the Anthropic system field plus messages array
// into the flat sdk.Message slice the CMaaS SDK expects (OpenAI shape).
//
// One Anthropic message can produce multiple sdk.Messages because Anthropic
// embeds tool_result blocks inside `user` messages, whereas OpenAI / the
// CMaaS protocol require a dedicated `role: "tool"` message per tool
// invocation. Order is preserved by flushing any pending text immediately
// before each tool_result expansion.
func buildSDKMessages(system json.RawMessage, in []anthropicMessage) ([]sdk.Message, error) {
	out := make([]sdk.Message, 0, len(in)+1)

	if hasJSONValue(system) {
		text, err := extractText(system, "system")
		if err != nil {
			return nil, err
		}
		out = append(out, sdk.Message{Role: "system", Content: text})
	}

	for i, m := range in {
		blocks, err := parseContentBlocks(m.Content, fmt.Sprintf("messages[%d].content", i))
		if err != nil {
			return nil, err
		}

		var pendingText strings.Builder
		var pendingToolCalls []sdk.ToolCall

		flushPending := func() {
			if pendingText.Len() == 0 && len(pendingToolCalls) == 0 {
				return
			}
			out = append(out, sdk.Message{
				Role:      m.Role,
				Content:   pendingText.String(),
				ToolCalls: pendingToolCalls,
			})
			pendingText.Reset()
			pendingToolCalls = nil
		}

		for j, b := range blocks {
			switch b.BlockType {
			case "text":
				pendingText.WriteString(b.Text)
			case "tool_use":
				if m.Role != "assistant" {
					return nil, fmt.Errorf("messages[%d].content[%d] tool_use is only valid in assistant messages", i, j)
				}
				pendingToolCalls = append(pendingToolCalls, sdk.ToolCall{
					ID:   b.ToolID,
					Type: "function",
					Function: sdk.FunctionCall{
						Name:      b.ToolName,
						Arguments: string(b.ToolInput),
					},
				})
			case "tool_result":
				if m.Role != "user" {
					return nil, fmt.Errorf("messages[%d].content[%d] tool_result is only valid in user messages", i, j)
				}
				flushPending()
				out = append(out, sdk.Message{
					Role:       "tool",
					ToolCallID: b.ToolUseID,
					Content:    b.ResultText,
				})
			}
		}
		flushPending()
	}
	return out, nil
}

func (h *handler) writeResponse(w http.ResponseWriter, completion *sdk.ChatCompletion, id, requestID, model string, toolChoiceNone bool) {
	toolCalls := ensureToolCallIDs(completion.ToolCalls)
	if toolChoiceNone {
		toolCalls = nil
	}

	content := make([]any, 0, 1+len(toolCalls))
	if completion.Content != "" {
		content = append(content, anthropicTextBlock{Type: "text", Text: completion.Content})
	}
	for _, tc := range toolCalls {
		var input map[string]any
		if tc.Function.Arguments != "" {
			// Tool arguments arrive as a JSON string; parse back into the
			// object shape Anthropic puts in tool_use.input.
			if err := json.Unmarshal([]byte(tc.Function.Arguments), &input); err != nil {
				log.Printf("[WARN] tool_call %s arguments are not valid JSON, surfacing empty input: %v", tc.ID, err)
			}
		}
		if input == nil {
			input = map[string]any{}
		}
		content = append(content, anthropicToolUseBlock{
			Type:  "tool_use",
			ID:    tc.ID,
			Name:  tc.Function.Name,
			Input: input,
		})
	}
	// When the model produced zero content, still emit a text block so the
	// content array isn't empty (matches Anthropic's behaviour on no-op
	// completions).
	if len(content) == 0 {
		content = append(content, anthropicTextBlock{Type: "text", Text: ""})
	}

	stopReason := "end_turn"
	if len(toolCalls) > 0 {
		stopReason = "tool_use"
	}

	resp := anthropicResponse{
		ID:           id,
		Type:         "message",
		Role:         "assistant",
		Content:      content,
		Model:        model,
		StopReason:   &stopReason,
		StopSequence: nil,
		Usage: anthropicUsage{
			InputTokens:  completion.Usage.PromptTokens,
			OutputTokens: completion.Usage.CompletionTokens,
		},
		ReasoningContent: normalizedReasoning(completion),
		RequestID:        requestID,
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

func normalizedReasoning(completion *sdk.ChatCompletion) string {
	if completion.ReasoningContent != "" {
		return completion.ReasoningContent
	}
	return completion.Reasoning
}

func (h *handler) writeStream(w http.ResponseWriter, completion *sdk.ChatCompletion, id, requestID, model string, toolChoiceNone bool) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeAnthropicError(w, http.StatusInternalServerError, "api_error", "streaming not supported", requestID)
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	writeEvent := func(name string, payload any) {
		data, _ := json.Marshal(payload)
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", name, data)
		flusher.Flush()
	}

	// Per-stream cache of generated IDs by tool-call index, so the
	// content_block_start carries a usable id even when the backend
	// returned empty on the first fragment.
	toolIDsByIndex := map[int]string{}
	ensureToolID := func(d sdk.ToolCallDelta) string {
		if d.ID != "" {
			toolIDsByIndex[d.Index] = d.ID
			return d.ID
		}
		if cached, ok := toolIDsByIndex[d.Index]; ok {
			return cached
		}
		generated := "toolu_" + uuid.New().String()
		toolIDsByIndex[d.Index] = generated
		return generated
	}

	// 1. message_start — usage filled in finally by message_delta.
	startMsg := anthropicResponse{
		ID:           id,
		Type:         "message",
		Role:         "assistant",
		Content:      []any{},
		Model:        model,
		StopReason:   nil, // null per Anthropic spec until message_delta finalizes it
		StopSequence: nil,
		Usage:        anthropicUsage{InputTokens: 0, OutputTokens: 0},
		RequestID:    requestID,
	}
	writeEvent("message_start", map[string]any{
		"type":    "message_start",
		"message": startMsg,
	})

	// Content-block state machine. Anthropic streams emit content blocks
	// strictly sequentially: content_block_start → deltas → content_block_stop
	// before the next block starts. We lazily allocate Anthropic indices as
	// either a text block (driven by completion.Chunks) or a tool_use block
	// (driven by completion.ToolCallChunks), closing the previous block
	// whenever the input source switches.
	currentBlockIdx := -1
	currentToolVLLMIdx := -1 // -1 means "current block is text (or none)"
	nextAnthIdx := 0
	sawToolCall := false

	closeCurrentBlock := func() {
		if currentBlockIdx == -1 {
			return
		}
		writeEvent("content_block_stop", map[string]any{
			"type":  "content_block_stop",
			"index": currentBlockIdx,
		})
		currentBlockIdx = -1
		currentToolVLLMIdx = -1
	}

	startTextBlock := func() {
		currentBlockIdx = nextAnthIdx
		currentToolVLLMIdx = -1
		nextAnthIdx++
		writeEvent("content_block_start", map[string]any{
			"type":          "content_block_start",
			"index":         currentBlockIdx,
			"content_block": anthropicTextBlock{Type: "text", Text: ""},
		})
	}

	startToolBlock := func(d sdk.ToolCallDelta) {
		currentBlockIdx = nextAnthIdx
		currentToolVLLMIdx = d.Index
		nextAnthIdx++
		writeEvent("content_block_start", map[string]any{
			"type":  "content_block_start",
			"index": currentBlockIdx,
			"content_block": map[string]any{
				"type":  "tool_use",
				"id":    ensureToolID(d),
				"name":  d.Name,
				"input": map[string]any{},
			},
		})
	}

	onText := func(text string) {
		if currentToolVLLMIdx != -1 {
			// Transition tool → text: close the tool block first.
			closeCurrentBlock()
		}
		if currentBlockIdx == -1 {
			startTextBlock()
		}
		writeEvent("content_block_delta", map[string]any{
			"type":  "content_block_delta",
			"index": currentBlockIdx,
			"delta": map[string]any{"type": "text_delta", "text": text},
		})
	}

	onToolDelta := func(d sdk.ToolCallDelta) {
		sawToolCall = true
		// Transition into a different block (text or a different tool).
		if currentBlockIdx != -1 && currentToolVLLMIdx != d.Index {
			closeCurrentBlock()
		}
		if currentBlockIdx == -1 {
			startToolBlock(d)
		}
		if d.Arguments != "" {
			writeEvent("content_block_delta", map[string]any{
				"type":  "content_block_delta",
				"index": currentBlockIdx,
				"delta": map[string]any{
					"type":         "input_json_delta",
					"partial_json": d.Arguments,
				},
			})
		}
	}

	// Multiplex text chunks and tool-call deltas.
	chunks := completion.Chunks
	toolChunks := completion.ToolCallChunks
	for chunks != nil || toolChunks != nil {
		select {
		case text, ok := <-chunks:
			if !ok {
				chunks = nil
				continue
			}
			if text == "" {
				continue
			}
			onText(text)
		case d, ok := <-toolChunks:
			if !ok {
				toolChunks = nil
				continue
			}
			if toolChoiceNone {
				continue // suppress tool calls when tool_choice is "none"
			}
			onToolDelta(d)
		}
	}

	if completion.Error != nil {
		log.Printf("[WARN] streaming error id=%s task_id=%s error=%v", id, completion.TaskID, completion.Error)
		closeCurrentBlock()
		writeEvent("error", map[string]any{
			"type": "error",
			"error": map[string]any{
				"type":    "api_error",
				"message": completion.Error.Error(),
			},
			"id":         id,
			"request_id": requestID,
		})
		return
	}

	closeCurrentBlock()

	stopReason := "end_turn"
	if sawToolCall {
		stopReason = "tool_use"
	}
	writeEvent("message_delta", map[string]any{
		"type": "message_delta",
		"delta": map[string]any{
			"stop_reason":   stopReason,
			"stop_sequence": nil,
		},
		"usage": map[string]any{
			"output_tokens": completion.Usage.CompletionTokens,
		},
	})

	writeEvent("message_stop", map[string]any{"type": "message_stop"})
}

// writeAnthropicError writes an error envelope. The id/request_id top-level
// fields mirror openai-proxy's writeOpenAIError so logs and responses use the
// same correlation id. Pass requestID="" to mint a fresh UUID.
func writeAnthropicError(w http.ResponseWriter, status int, errType, msg, requestID string) {
	if requestID == "" {
		requestID = uuid.New().String()
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(map[string]any{
		"type":       "error",
		"error":      map[string]any{"type": errType, "message": msg},
		"id":         "msg_" + requestID,
		"request_id": requestID,
	})
}

// ---- main ----

func main() {
	endpoint := flag.String("endpoint", "", "CMaaS WebSocket endpoint (required)")
	model := flag.String("model", "", "Model name (required)")
	listen := flag.String("listen", ":8081", "HTTP listen address")
	workers := flag.Int("workers", 8, "Number of concurrent CMaaS client sessions (= max concurrency)")
	cacheTTL := flag.Duration("cache-ttl", 12*time.Hour, "Public key cache TTL (set 0 to disable)")
	flag.Parse()

	if *endpoint == "" {
		log.Fatal("endpoint flag is required")
	}
	if *model == "" {
		log.Fatal("model flag is required")
	}

	apiKey := os.Getenv("DASHSCOPE_API_KEY")
	if apiKey == "" {
		log.Fatal("DASHSCOPE_API_KEY environment variable is required")
	}

	ttl := *cacheTTL
	if ttl == 0 {
		ttl = -1
	}

	opts := &sdk.Options{
		Endpoint:       *endpoint,
		APIKey:         apiKey,
		Model:          *model,
		ReuseSession:   false,
		PubKeyCacheTTL: ttl,
	}

	log.Printf("Creating pool of %d CMaaS client sessions...", *workers)
	p, err := newPool(*workers, opts)
	if err != nil {
		log.Fatalf("Failed to create client pool: %v", err)
	}
	defer p.close()

	h := &handler{pool: p, model: *model}

	srv := &http.Server{
		Addr:         *listen,
		Handler:      h,
		ReadTimeout:  5 * time.Minute,
		WriteTimeout: 10 * time.Minute,
		IdleTimeout:  2 * time.Minute,
	}

	log.Printf("Anthropic-compatible proxy listening on %s (workers=%d)", *listen, *workers)
	log.Printf("CMaaS backend endpoint=%s model=%q reuse_session=false", *endpoint, *model)
	log.Printf("  POST /v1/messages")
	log.Printf("  GET  /v1/models")
	log.Printf("")
	log.Printf("Example (Anthropic Python SDK):")
	log.Printf("  client = anthropic.Anthropic(base_url=\"http://localhost%s\", api_key=\"ignored\")", *listen)
	log.Printf("  client.messages.create(model=%q, max_tokens=64,", *model)
	log.Printf("      messages=[{\"role\":\"user\",\"content\":\"hello\"}])")
	log.Printf("")

	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("Server error: %v", err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	srv.Shutdown(ctx)
}
