package agent

import (
	"encoding/json"
	"strings"
)

// ---------------------------------------------------------------------------
// Source-faithful Pi RPC wire types.
//
// These structs mirror the current typed unions from pi-mono commit
// eb1f87fa9a29e27e0c63dcb40dbed9a3624c82b1 (v0.84.2):
//   - packages/agent/src/types.ts  (AgentEvent, AgentMessage, AssistantMessageEvent,
//     ToolResultMessage, Usage, StopReason)
//   - packages/coding-agent/src/core/agent-session.ts (AgentSessionEvent union)
//   - packages/coding-agent/src/modes/json-event.ts  (ToJsonEvent strips the
//     cumulative message snapshot from message_update, keeping usage + delta)
//   - packages/ai/src/types.ts (AssistantMessage, ToolCall, TextContent,
//     ThinkingContent, ImageContent)
//
// Tool-execution, queue, retry, compaction, bash and agent_end fields are all
// TOP-LEVEL frame fields, not nested under a `data` envelope.
// ---------------------------------------------------------------------------

// piRPCFrame is the union of every stdout frame (response or event) and every
// stdin command. Frame classification: a `response` type with an `id` is a
// command response; otherwise the frame is an event.
//
// Several JSON keys are reused across event types with type-specific meaning
// (e.g. `result` on tool_execution_end vs compaction_end; `success` on a
// response vs auto_retry_end). Because only one event type is present per
// frame, each JSON key maps to a single field and the handler interprets it
// based on the frame's `type`.
type piRPCFrame struct {
	// --- common ---
	Type string          `json:"type"`
	ID   json.RawMessage `json:"id,omitempty"`

	// --- command response ---
	Command string          `json:"command,omitempty"`
	Success *bool           `json:"success,omitempty"`
	Error   string          `json:"error,omitempty"`
	Data    json.RawMessage `json:"data,omitempty"`

	// --- streaming events ---
	// message_update: latest cumulative provider usage + the delta event.
	Usage                 json.RawMessage `json:"usage,omitempty"`
	AssistantMessageEvent json.RawMessage `json:"assistantMessageEvent,omitempty"`

	// message_start / message_end / turn_end.message: a full AgentMessage.
	Message *piAgentMessage `json:"message,omitempty"`

	// turn_end.toolResults: ToolResultMessage[].
	ToolResults []piAgentMessage `json:"toolResults,omitempty"`

	// agent_end.messages: AgentMessage[] + willRetry.
	Messages  []piAgentMessage `json:"messages,omitempty"`
	WillRetry bool             `json:"willRetry,omitempty"`

	// tool_execution_start/update/end (top-level).
	ToolCallID    string          `json:"toolCallId,omitempty"`
	ToolName      string          `json:"toolName,omitempty"`
	Args          json.RawMessage `json:"args,omitempty"`
	PartialResult json.RawMessage `json:"partialResult,omitempty"`
	Result        json.RawMessage `json:"result,omitempty"`
	IsError       bool            `json:"isError,omitempty"`

	// queue_update.
	Steering []string `json:"steering,omitempty"`
	FollowUp []string `json:"followUp,omitempty"`

	// compaction_start / compaction_end.
	Reason  string `json:"reason,omitempty"`
	Aborted bool   `json:"aborted,omitempty"`

	// auto_retry_start / auto_retry_end.
	Attempt      int    `json:"attempt,omitempty"`
	MaxAttempts  int    `json:"maxAttempts,omitempty"`
	DelayMs      int    `json:"delayMs,omitempty"`
	ErrorMessage string `json:"errorMessage,omitempty"`
	FinalError   string `json:"finalError,omitempty"`

	// bash_execution_update.
	Delta string `json:"delta,omitempty"`

	// extension_error.
	ExtensionPath  string `json:"extensionPath,omitempty"`
	ExtensionEvent string `json:"event,omitempty"`
}

// IDString returns the raw id field as a string (for bash_execution_update
// correlation), stripping surrounding quotes.
func (f piRPCFrame) IDString() string {
	if len(f.ID) == 0 {
		return ""
	}
	var s string
	if json.Unmarshal(f.ID, &s) == nil {
		return s
	}
	return strings.Trim(string(f.ID), `"`)
}

// CompactionResultPresent reports whether a compaction_end frame carries a
// non-null result payload (a genuinely successful compaction reports a
// CompactionResult; an aborted/failed completion carries result:null).
func (f piRPCFrame) CompactionResultPresent() bool {
	// result:null serializes as the JSON "null"; treat it as absent.
	if len(f.Result) == 0 {
		return false
	}
	return strings.TrimSpace(string(f.Result)) != "null"
}

// CompactionTooSmallNoop identifies Pi's benign short-session compact rejection.
func (f piRPCFrame) CompactionTooSmallNoop() bool {
	message := strings.ToLower(strings.TrimSpace(f.ErrorMessage))
	return strings.Contains(message, "nothing to compact") && strings.Contains(message, "session too small")
}

// piAgentMessage is the union of UserMessage | AssistantMessage |
// ToolResultMessage | BashExecutionMessage | CustomMessageEntry as serialized
// by pi. Field presence is governed by the message `role`.
type piAgentMessage struct {
	Role string `json:"role"`

	// content: string OR (TextContent | ThinkingContent | ToolCall | ImageContent)[].
	Content json.RawMessage `json:"content,omitempty"`

	// AssistantMessage fields.
	API           string   `json:"api,omitempty"`
	Provider      string   `json:"provider,omitempty"`
	Model         string   `json:"model,omitempty"`
	ResponseModel string   `json:"responseModel,omitempty"`
	ResponseID    string   `json:"responseId,omitempty"`
	Usage         *piUsage `json:"usage,omitempty"`
	StopReason    string   `json:"stopReason,omitempty"`
	RawStopReason string   `json:"rawStopReason,omitempty"`
	EndTurn       *bool    `json:"endTurn,omitempty"`
	ErrorMessage  string   `json:"errorMessage,omitempty"`
	Timestamp     int64    `json:"timestamp,omitempty"`

	// ToolResultMessage fields.
	ToolCallID string `json:"toolCallId,omitempty"`
	ToolName   string `json:"toolName,omitempty"`
	IsError    bool   `json:"isError,omitempty"`

	// BashExecutionMessage fields.
	Command        string `json:"command,omitempty"`
	Output         string `json:"output,omitempty"`
	ExitCode       *int   `json:"exitCode,omitempty"`
	Cancelled      bool   `json:"cancelled,omitempty"`
	Truncated      bool   `json:"truncated,omitempty"`
	FullOutputPath string `json:"fullOutputPath,omitempty"`
}

// piContentBlock mirrors TextContent | ThinkingContent | ToolCall | ImageContent.
type piContentBlock struct {
	Type string `json:"type"`

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

	// thinking
	Thinking          string `json:"thinking,omitempty"`
	ThinkingSignature string `json:"thinkingSignature,omitempty"`
	Redacted          bool   `json:"redacted,omitempty"`

	// toolCall
	ID        string          `json:"id,omitempty"`
	Name      string          `json:"name,omitempty"`
	Arguments json.RawMessage `json:"arguments,omitempty"`

	// image
	Data     string `json:"data,omitempty"`
	MimeType string `json:"mimeType,omitempty"`
}

// piUsage mirrors the pi-ai Usage type.
type piUsage struct {
	Input       int  `json:"input"`
	Output      int  `json:"output"`
	CacheRead   int  `json:"cacheRead"`
	CacheWrite  int  `json:"cacheWrite"`
	Reasoning   *int `json:"reasoning,omitempty"`
	TotalTokens int  `json:"totalTokens"`
}

// piAssistantMessageEvent is the wire form of AssistantMessageEvent WITHOUT the
// cumulative `partial` snapshot (json-event.ts strips it). Only the delta fields
// relevant to a client remain: type, contentIndex, delta|content|toolCall, and
// for the terminal done/error variants the reason + authoritative message.
type piAssistantMessageEvent struct {
	Type         string          `json:"type"`
	ContentIndex int             `json:"contentIndex,omitempty"`
	Delta        string          `json:"delta,omitempty"`
	Content      string          `json:"content,omitempty"`
	ToolCall     *piToolCallWire `json:"toolCall,omitempty"`
	Reason       string          `json:"reason,omitempty"`
	Message      json.RawMessage `json:"message,omitempty"`
}

// piToolCallWire mirrors the ToolCall content block used by toolcall_end.
type piToolCallWire struct {
	ID        string          `json:"id"`
	Name      string          `json:"name"`
	Arguments json.RawMessage `json:"arguments"`
}

// piToolExecutionPartialResult mirrors the accumulated partial tool result.
// partialResult is CUMULATIVE (clients replace their display, not append).
type piToolExecutionPartialResult struct {
	Content []piContentBlock `json:"content"`
	Details json.RawMessage  `json:"details,omitempty"`
}

// piToolExecutionResult mirrors the final tool result on tool_execution_end.
type piToolExecutionResult struct {
	Content []piContentBlock `json:"content"`
	Details json.RawMessage  `json:"details,omitempty"`
}