// Package llm is a hand-rolled client for the Anthropic Messages API.
//
// The official anthropic-sdk-go would cover most of this; the types are written
// out here because understanding the wire format is the point of the project,
// and Phase 5 hand-accumulates streaming tool-call JSON.
package llm

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
)

// Roles and block types, spelled out rather than inlined as string literals.
const (
	RoleUser      = "user"
	RoleAssistant = "assistant"

	BlockText             = "text"
	BlockImage            = "image"
	BlockThinking         = "thinking"
	BlockRedactedThinking = "redacted_thinking"
	BlockToolUse          = "tool_use"
	BlockToolResult       = "tool_result"

	// BlockServerToolUse and BlockWebSearchToolResult are the pair a hosted
	// web search produces. The API runs the tool and supplies both blocks;
	// Codey only has to show them and send them back unchanged. They are kept
	// distinct from tool_use/tool_result so nothing local ever dispatches one.
	BlockServerToolUse       = "server_tool_use"
	BlockWebSearchToolResult = "web_search_tool_result"
)

// WebSearchToolType and WebSearchToolName identify Anthropic's hosted web
// search in the tools array. The type carries a date because the API versions
// hosted tools by it; the name is fixed by the API and is what comes back in
// server_tool_use.name.
//
// This is deliberately the oldest version rather than the newest. Later ones
// (web_search_20260209 onwards) add dynamic filtering, where search results are
// reduced by code the model writes before they reach the context window — worth
// having, and gated on Opus 4.6 and later. Codey sends whatever model id the
// user names, including ones only a gateway knows about, so pinning the newest
// here would 400 the entire request for anyone on an older model. Choosing the
// version from the model in force is the fix; until that exists, this is the
// one that works everywhere the hosted tool works at all.
const (
	WebSearchToolType = "web_search_20250305"
	WebSearchToolName = "web_search"
)

// Stop reasons the loop has to distinguish.
const (
	StopEndTurn   = "end_turn"
	StopMaxTokens = "max_tokens"
	StopToolUse   = "tool_use"
	StopPauseTurn = "pause_turn"
	StopRefusal   = "refusal"
)

// Message is one turn of conversation.
type Message struct {
	Role    string         `json:"role"`
	Content []ContentBlock `json:"content"`
}

// ContentBlock is the flattened union of every block type Codey handles.
// Only the fields belonging to Type are serialized, so blocks survive the
// round-trip back to the API — which matters most for thinking blocks, whose
// signature must be echoed back unchanged.
type ContentBlock struct {
	Type string `json:"type"`

	// text
	Text string `json:"text,omitempty"`

	// image
	Source *ImageSource `json:"source,omitempty"`

	// thinking / redacted_thinking. Thinking is deliberately not omitempty:
	// with display omitted (the default on Opus 5) the text is empty but the
	// field must still be sent back.
	Thinking  string `json:"thinking"`
	Signature string `json:"signature,omitempty"`
	Data      string `json:"data,omitempty"`

	// tool_use, and server_tool_use for a hosted tool the API ran itself.
	ID    string          `json:"id,omitempty"`
	Name  string          `json:"name,omitempty"`
	Input json.RawMessage `json:"input,omitempty"`

	// tool_result. Content is the text of the result — what every reader
	// that only understands text gets. Blocks is set when the result is
	// more than text can carry, an image above all: then the API is sent the
	// blocks and Content holds their text so the transcript, the log and the
	// staleness check still have something to show. A web_search_tool_result
	// carries an array of a different shape here, which is left to raw (see
	// UnmarshalJSON) rather than modelled.
	ToolUseID string         `json:"tool_use_id,omitempty"`
	Content   string         `json:"content,omitempty"`
	Blocks    []ContentBlock `json:"-"`
	IsError   bool           `json:"is_error,omitempty"`

	// CacheControl marks this block as a cache breakpoint.
	CacheControl *CacheControl `json:"cache_control,omitempty"`

	// raw holds the original bytes of a block type this struct does not model,
	// so an unfamiliar block still round-trips instead of being flattened away.
	raw json.RawMessage
}

func TextBlock(s string) ContentBlock { return ContentBlock{Type: BlockText, Text: s} }

// ImageSource is where the bytes of an image block come from. Only base64 is
// modelled: a URL source would have the API fetch the image itself, which is
// the wrong shape for pictures pasted off a local clipboard.
type ImageSource struct {
	Type      string `json:"type"`
	MediaType string `json:"media_type"`
	Data      string `json:"data"`
}

// ImageBlock wraps already-encoded base64 bytes. mediaType is one of the four
// the API accepts — image/png, image/jpeg, image/gif, image/webp.
func ImageBlock(mediaType, b64 string) ContentBlock {
	return ContentBlock{Type: BlockImage, Source: &ImageSource{Type: "base64", MediaType: mediaType, Data: b64}}
}

// ToolResultBlock builds the reply to one tool_use. Every tool_use id must get
// one of these, including failures — a missing result is an API error.
func ToolResultBlock(toolUseID, content string, isError bool) ContentBlock {
	if content == "" {
		// An empty content field is rejected; say something instead.
		content = "(no output)"
	}
	return ContentBlock{Type: BlockToolResult, ToolUseID: toolUseID, Content: content, IsError: isError}
}

// ToolResultBlocks builds a successful reply whose content is blocks rather
// than a string — a caption and an image, typically. Content is filled with
// the text blocks so the transcript and the log have words where the model
// has a picture. With no blocks it is an ordinary empty result.
func ToolResultBlocks(toolUseID string, blocks []ContentBlock) ContentBlock {
	if len(blocks) == 0 {
		return ToolResultBlock(toolUseID, "", false)
	}
	var text []string
	for _, b := range blocks {
		if b.Type == BlockText && b.Text != "" {
			text = append(text, b.Text)
		}
	}
	return ContentBlock{Type: BlockToolResult, ToolUseID: toolUseID, Content: strings.Join(text, "\n"), Blocks: blocks}
}

func (b *ContentBlock) UnmarshalJSON(data []byte) error {
	type alias ContentBlock
	var a alias
	if err := json.Unmarshal(data, &a); err != nil {
		// A web_search_tool_result carries "content" as an array of results
		// or an error object, not the string this struct models. Rejecting it
		// would take the whole response down with it, and the bytes are kept
		// in raw for the round trip anyway, so a type mismatch on that one
		// field is tolerated. Anything else is still a malformed block.
		var typeErr *json.UnmarshalTypeError
		if !errors.As(err, &typeErr) || typeErr.Field != "content" {
			return err
		}
	}
	*b = ContentBlock(a)
	b.raw = append(json.RawMessage(nil), data...)
	if b.Type == BlockToolResult {
		// Our own tool_result with block content, read back from a saved
		// session. Recover the blocks and the caption; only a well-formed
		// array counts, anything else stays as it was in raw.
		var probe struct {
			Content json.RawMessage `json:"content"`
		}
		if json.Unmarshal(data, &probe) == nil && len(probe.Content) > 0 && probe.Content[0] == '[' {
			var blocks []ContentBlock
			if json.Unmarshal(probe.Content, &blocks) == nil {
				rebuilt := ToolResultBlocks(b.ToolUseID, blocks)
				b.Blocks, b.Content = rebuilt.Blocks, rebuilt.Content
			}
		}
	}
	return nil
}

// Raw returns the block exactly as it arrived from the API, or nil for a block
// built locally. It exists for block types this struct does not model field by
// field — a web_search_tool_result's content, above all — so a caller who wants
// to look inside one can, without the struct growing a field per hosted tool.
func (b ContentBlock) Raw() json.RawMessage { return b.raw }

// MarshalJSON emits exactly the fields that belong to Type. Marshaling the
// whole struct would send stray keys (an empty "thinking" on every text block)
// and the API rejects some of those.
func (b ContentBlock) MarshalJSON() ([]byte, error) {
	m := map[string]any{"type": b.Type}
	switch b.Type {
	case BlockText:
		m["text"] = b.Text
	case BlockImage:
		if b.Source == nil {
			return nil, fmt.Errorf("llm: image block has no source")
		}
		m["source"] = b.Source
	case BlockThinking:
		m["thinking"] = b.Thinking
		if b.Signature != "" {
			m["signature"] = b.Signature
		}
	case BlockRedactedThinking:
		m["data"] = b.Data
	case BlockToolUse, BlockServerToolUse:
		// server_tool_use is deliberately not left to the raw path below.
		// When streamed, raw is the content_block_start snapshot in which
		// "input" is still {}; the query arrives afterwards as input_json_delta
		// fragments and lands in Input. Re-sending the snapshot after
		// pause_turn would hand the API a search with no query.
		m["id"] = b.ID
		m["name"] = b.Name
		if len(b.Input) == 0 {
			m["input"] = json.RawMessage("{}")
		} else {
			m["input"] = b.Input
		}
	case BlockToolResult:
		m["tool_use_id"] = b.ToolUseID
		if len(b.Blocks) > 0 {
			m["content"] = b.Blocks
		} else {
			m["content"] = b.Content
		}
		if b.IsError {
			m["is_error"] = true
		}
	default:
		// Everything else — web_search_tool_result included — goes back as it
		// came. That block arrives whole, with no deltas, so raw is complete;
		// and its content is the API's, encrypted, and not ours to reshape.
		if len(b.raw) == 0 {
			return nil, fmt.Errorf("llm: cannot marshal unknown content block %q", b.Type)
		}
		if b.CacheControl == nil {
			return b.raw, nil
		}
		return withCacheControl(b.raw, b.CacheControl)
	}
	if b.CacheControl != nil {
		m["cache_control"] = b.CacheControl
	}
	return json.Marshal(m)
}

// withCacheControl attaches a cache breakpoint to a raw block. The loop puts
// its rolling breakpoint on the last block of the last message, and after a
// pause_turn that block is usually a web_search_tool_result — returning raw
// unchanged there would silently drop the breakpoint and miss the cache on
// every search turn. Values are kept as json.RawMessage so nothing inside the
// block (numbers, encrypted payloads) is reinterpreted on the way through.
func withCacheControl(raw json.RawMessage, cc *CacheControl) ([]byte, error) {
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil {
		return nil, fmt.Errorf("llm: raw content block is not an object: %w", err)
	}
	ccJSON, err := json.Marshal(cc)
	if err != nil {
		return nil, err
	}
	obj["cache_control"] = ccJSON
	return json.Marshal(obj)
}

// SystemBlock is a system-prompt segment. It is a list rather than a string so
// a cache breakpoint can be attached in Phase 4.
type SystemBlock struct {
	Type         string        `json:"type"`
	Text         string        `json:"text"`
	CacheControl *CacheControl `json:"cache_control,omitempty"`
}

type CacheControl struct {
	Type string `json:"type"`
	TTL  string `json:"ttl,omitempty"`
}

func Ephemeral() *CacheControl { return &CacheControl{Type: "ephemeral"} }

// ToolDef is one entry of the request's tools array. It covers two shapes: a
// local tool, declared by name, description and input_schema; and a hosted
// tool, declared by type and name alone, which the API runs itself.
//
// Description and InputSchema are omitempty because of the second shape. A
// hosted tool has no schema — the API owns it — and a non-omitempty
// InputSchema would send "input_schema":null, which is a 400. Local tools
// always fill both, so they lose nothing.
type ToolDef struct {
	// Type is set only for a hosted tool, e.g. WebSearchToolType. A local
	// tool leaves it empty and the API infers the custom-tool shape.
	Type        string          `json:"type,omitempty"`
	Name        string          `json:"name"`
	Description string          `json:"description,omitempty"`
	InputSchema json.RawMessage `json:"input_schema,omitempty"`

	// The remaining fields are hosted web search knobs. MaxUses caps the
	// searches in one response; zero means the API default rather than none,
	// hence omitempty. AllowedDomains and BlockedDomains are mutually
	// exclusive on the API side, and it is a 400 to send both.
	MaxUses        int      `json:"max_uses,omitempty"`
	AllowedDomains []string `json:"allowed_domains,omitempty"`
	BlockedDomains []string `json:"blocked_domains,omitempty"`
}

// WebSearchTool is the tools-array entry for Anthropic's hosted web search.
// maxUses <= 0 leaves the cap to the API default. The definition is constant
// for a session, which is what lets it sit in the cached prompt prefix.
func WebSearchTool(maxUses int) ToolDef {
	def := ToolDef{Type: WebSearchToolType, Name: WebSearchToolName}
	if maxUses > 0 {
		def.MaxUses = maxUses
	}
	return def
}

// Thinking configures extended thinking. Adaptive is the only mode current
// models accept; budget_tokens is a 400 on Opus 5.
type Thinking struct {
	Type    string `json:"type"`
	Display string `json:"display,omitempty"`
}

func Adaptive() *Thinking { return &Thinking{Type: "adaptive"} }

// OutputConfig carries the effort knob. Omitted means the default, "high".
type OutputConfig struct {
	Effort string `json:"effort,omitempty"`
}

// EffortLevels are the values output_config.effort accepts, cheapest first.
// Anything else is a 400 from the API, so callers check before sending.
var EffortLevels = []string{"low", "medium", "high", "xhigh", "max"}

// ValidEffort reports whether level is one the API accepts. The empty string
// is not: it means "say nothing", which callers handle by omitting the field.
func ValidEffort(level string) bool {
	for _, l := range EffortLevels {
		if l == level {
			return true
		}
	}
	return false
}

// Request is the body of POST /v1/messages.
type Request struct {
	Model        string        `json:"model"`
	MaxTokens    int           `json:"max_tokens"`
	System       []SystemBlock `json:"system,omitempty"`
	Messages     []Message     `json:"messages"`
	Tools        []ToolDef     `json:"tools,omitempty"`
	Thinking     *Thinking     `json:"thinking,omitempty"`
	OutputConfig *OutputConfig `json:"output_config,omitempty"`
	Stream       bool          `json:"stream,omitempty"`
}

// Usage is the token accounting returned with every response.
type Usage struct {
	InputTokens              int `json:"input_tokens"`
	OutputTokens             int `json:"output_tokens"`
	CacheReadInputTokens     int `json:"cache_read_input_tokens"`
	CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
}

// StopDetails is populated only when StopReason is "refusal".
type StopDetails struct {
	Type        string `json:"type"`
	Category    string `json:"category"`
	Explanation string `json:"explanation"`
}

// Response is one assistant turn.
type Response struct {
	ID          string         `json:"id"`
	Type        string         `json:"type"`
	Role        string         `json:"role"`
	Model       string         `json:"model"`
	Content     []ContentBlock `json:"content"`
	StopReason  string         `json:"stop_reason"`
	StopDetails *StopDetails   `json:"stop_details"`
	Usage       Usage          `json:"usage"`
}

// Message converts the response into a history entry.
func (r *Response) Message() Message {
	return Message{Role: RoleAssistant, Content: r.Content}
}

// ToolUses returns the local tool calls in this response, in order. It matches
// tool_use alone, never server_tool_use: everything it returns is dispatched
// to the local registry and owed a tool_result, and a hosted call is neither.
func (r *Response) ToolUses() []ContentBlock {
	var out []ContentBlock
	for _, b := range r.Content {
		if b.Type == BlockToolUse {
			out = append(out, b)
		}
	}
	return out
}

// ServerToolUses returns the hosted tool calls the API ran on this turn, in
// order. They are for showing, not dispatching: the API has already executed
// them and put each result in the same response.
func (r *Response) ServerToolUses() []ContentBlock {
	var out []ContentBlock
	for _, b := range r.Content {
		if b.Type == BlockServerToolUse {
			out = append(out, b)
		}
	}
	return out
}

// Text concatenates the response's text blocks.
func (r *Response) Text() string {
	var s string
	for _, b := range r.Content {
		if b.Type == BlockText {
			s += b.Text
		}
	}
	return s
}