// 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"
	"fmt"
)

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

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

// 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 the harness 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"`

	// 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
	ID    string          `json:"id,omitempty"`
	Name  string          `json:"name,omitempty"`
	Input json.RawMessage `json:"input,omitempty"`

	// tool_result
	ToolUseID string `json:"tool_use_id,omitempty"`
	Content   string `json:"content,omitempty"`
	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} }

// 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}
}

func (b *ContentBlock) UnmarshalJSON(data []byte) error {
	type alias ContentBlock
	var a alias
	if err := json.Unmarshal(data, &a); err != nil {
		return err
	}
	*b = ContentBlock(a)
	b.raw = append(json.RawMessage(nil), data...)
	return nil
}

// 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 BlockThinking:
		m["thinking"] = b.Thinking
		if b.Signature != "" {
			m["signature"] = b.Signature
		}
	case BlockRedactedThinking:
		m["data"] = b.Data
	case BlockToolUse:
		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
		m["content"] = b.Content
		if b.IsError {
			m["is_error"] = true
		}
	default:
		if len(b.raw) > 0 {
			return b.raw, nil
		}
		return nil, fmt.Errorf("llm: cannot marshal unknown content block %q", b.Type)
	}
	if b.CacheControl != nil {
		m["cache_control"] = b.CacheControl
	}
	return json.Marshal(m)
}

// 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.
type ToolDef struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"input_schema"`
}

// 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"`
}

// 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 tool calls in this response, in order.
func (r *Response) ToolUses() []ContentBlock {
	var out []ContentBlock
	for _, b := range r.Content {
		if b.Type == BlockToolUse {
			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
}