package agent

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

type appServerEnvelope struct {
	ID     json.RawMessage         `json:"id,omitempty"`
	Method string                  `json:"method,omitempty"`
	Params json.RawMessage         `json:"params,omitempty"`
	Result json.RawMessage         `json:"result,omitempty"`
	Error  *appServerResponseError `json:"error,omitempty"`
}

type appServerResponseError struct {
	Code    int             `json:"code"`
	Message string          `json:"message"`
	Data    json.RawMessage `json:"data,omitempty"`
}

type appServerRequestError struct {
	Method  string
	Code    int
	Message string
	Data    json.RawMessage
	Cause   error
}

type appServerMCPServerToolCallParams struct {
	Arguments json.RawMessage `json:"arguments,omitempty"`
	CallID    string          `json:"callId,omitempty"`
	Server    string          `json:"server"`
	ThreadID  string          `json:"threadId"`
	Tool      string          `json:"tool"`
	TurnID    string          `json:"turnId,omitempty"`
}

type appServerDynamicTool struct {
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"inputSchema"`
	Name        string          `json:"name"`
}

type appServerMCPToolsListResult struct {
	Tools []appServerMCPToolDef `json:"tools"`
}

type appServerMCPToolDef struct {
	Description string                   `json:"description"`
	InputSchema json.RawMessage          `json:"inputSchema"`
	Name        string                   `json:"name"`
	Annotations *types.MCPToolAnnotation `json:"annotations,omitempty"`
}

type appServerDynamicToolCallParams struct {
	Arguments json.RawMessage `json:"arguments,omitempty"`
	CallID    string          `json:"callId"`
	Namespace string          `json:"namespace,omitempty"`
	ThreadID  string          `json:"threadId"`
	Tool      string          `json:"tool"`
	TurnID    string          `json:"turnId"`
}

type appServerMCPToolCallResult struct {
	Content           []types.MCPToolContent `json:"content"`
	StructuredContent any                    `json:"structuredContent,omitempty"`

	runtimeDeliveryReceipts []string
	runtimeDeliverySession  string
}

type appServerDynamicToolCallOutputContentItem struct {
	Text string `json:"text,omitempty"`
	Type string `json:"type"`
}

type appServerDynamicToolCallResponse struct {
	ContentItems []appServerDynamicToolCallOutputContentItem `json:"contentItems"`
	Success      bool                                        `json:"success"`

	runtimeDeliveryReceipts []string
	runtimeDeliverySession  string
}

type appServerMCPServerElicitationRequestParams struct {
	Message         string                                       `json:"message"`
	Meta            map[string]any                               `json:"_meta,omitempty"`
	Mode            string                                       `json:"mode"`
	RequestedSchema appServerMCPServerElicitationRequestedSchema `json:"requestedSchema"`
	ServerName      string                                       `json:"serverName"`
	ThreadID        string                                       `json:"threadId"`
	TurnID          string                                       `json:"turnId,omitempty"`
}

type appServerMCPServerElicitationRequestedSchema struct {
	Properties map[string]appServerMCPServerElicitationProperty `json:"properties"`
	Required   []string                                         `json:"required,omitempty"`
	Type       string                                           `json:"type"`
}

type appServerMCPServerElicitationProperty struct {
	AnyOf   []appServerMCPServerElicitationConstOption `json:"anyOf,omitempty"`
	Default json.RawMessage                            `json:"default,omitempty"`
	Enum    []string                                   `json:"enum,omitempty"`
	OneOf   []appServerMCPServerElicitationConstOption `json:"oneOf,omitempty"`
	Type    string                                     `json:"type"`
}

type appServerMCPServerElicitationConstOption struct {
	Const string `json:"const"`
	Title string `json:"title,omitempty"`
}

type appServerMCPServerElicitationRequestResponse struct {
	Action  string         `json:"action"`
	Content map[string]any `json:"content,omitempty"`
	Meta    map[string]any `json:"_meta,omitempty"`
}

type appServerPermissionsRequestApprovalParams struct {
	ItemID      string         `json:"itemId"`
	Permissions map[string]any `json:"permissions"`
	Reason      string         `json:"reason,omitempty"`
	ThreadID    string         `json:"threadId"`
	TurnID      string         `json:"turnId"`
}

type appServerPermissionsRequestApprovalResponse struct {
	Permissions map[string]any `json:"permissions"`
	Scope       string         `json:"scope,omitempty"`
}

type appServerServerRequestResolvedNotification struct {
	RequestID json.RawMessage `json:"requestId"`
	ThreadID  string          `json:"threadId"`
}

func (e *appServerRequestError) Error() string {
	message := strings.TrimSpace(e.Message)
	if message == "" {
		message = fmt.Sprintf("JSON-RPC error %d", e.Code)
	}
	method := strings.TrimSpace(e.Method)
	if method == "" {
		return message
	}
	return fmt.Sprintf("%s: %s", method, message)
}

func (e *appServerRequestError) Unwrap() error {
	if e != nil && e.Code == codexAppServerTransportFailureCode {
		if e.Cause != nil {
			return errors.Join(errCodexAppServerTransportAmbiguous, e.Cause)
		}
		return errCodexAppServerTransportAmbiguous
	}
	if e != nil {
		return e.Cause
	}
	return nil
}

// ---------------------------------------------------------------------------
// Thread & turn types
// ---------------------------------------------------------------------------

type appServerTurnResult struct {
	Turn appServerTurn `json:"turn"`
}

type appServerTurnIDResult struct {
	TurnID string `json:"turnId"`
}

type appServerThreadResult struct {
	Thread        appServerThread `json:"thread"`
	Model         string          `json:"model"`
	ModelProvider string          `json:"modelProvider"`
}

type appServerThreadStatus struct {
	Type        string   `json:"type"`
	ActiveFlags []string `json:"activeFlags,omitempty"`
}

func (status *appServerThreadStatus) UnmarshalJSON(data []byte) error {
	type rawStatus appServerThreadStatus

	statusType, ok, err := unmarshalAppServerStatusType(data)
	if err != nil {
		return err
	}
	if ok {
		*status = appServerThreadStatus{Type: statusType}
		return nil
	}

	var decoded rawStatus
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	*status = appServerThreadStatus(decoded)
	return nil
}

type appServerThread struct {
	ID     string                `json:"id"`
	Status appServerThreadStatus `json:"status"`
	Turns  []appServerTurn       `json:"turns"`
}

type appServerTurnStatus struct {
	Type string `json:"type"`
}

func (status *appServerTurnStatus) UnmarshalJSON(data []byte) error {
	type rawStatus appServerTurnStatus

	statusType, ok, err := unmarshalAppServerStatusType(data)
	if err != nil {
		return err
	}
	if ok {
		*status = appServerTurnStatus{Type: statusType}
		return nil
	}

	var decoded rawStatus
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	*status = appServerTurnStatus(decoded)
	return nil
}

type appServerTurn struct {
	ID        string              `json:"id"`
	Items     []appServerItem     `json:"items,omitempty"`
	ItemsView string              `json:"itemsView,omitempty"`
	Status    appServerTurnStatus `json:"status"`
}

// ---------------------------------------------------------------------------
// Notification types
// ---------------------------------------------------------------------------

type appServerTurnNotification struct {
	ThreadID string        `json:"threadId"`
	Turn     appServerTurn `json:"turn"`
}

type appServerThreadStatusNotification struct {
	ThreadID string                `json:"threadId"`
	Status   appServerThreadStatus `json:"status"`
}

type appServerThreadSettingsUpdatedNotification struct {
	ThreadID       string `json:"threadId"`
	ThreadSettings struct {
		Model string `json:"model"`
	} `json:"threadSettings"`
}

type appServerModelReroutedNotification struct {
	ThreadID string `json:"threadId"`
	TurnID   string `json:"turnId"`
	ToModel  string `json:"toModel"`
}

type appServerThreadTokenUsage struct {
	Last               appServerTokenUsageBreakdown `json:"last"`
	ModelContextWindow *int64                       `json:"modelContextWindow"`
	Total              appServerTokenUsageBreakdown `json:"total"`
}

type appServerTokenUsageBreakdown struct {
	CachedInputTokens     int64 `json:"cachedInputTokens"`
	InputTokens           int64 `json:"inputTokens"`
	OutputTokens          int64 `json:"outputTokens"`
	ReasoningOutputTokens int64 `json:"reasoningOutputTokens"`
	TotalTokens           int64 `json:"totalTokens"`
}

type appServerThreadTokenUsageNotification struct {
	ThreadID   string                    `json:"threadId"`
	TokenUsage appServerThreadTokenUsage `json:"tokenUsage"`
	TurnID     string                    `json:"turnId"`
}

type appServerItemNotification struct {
	ThreadID string        `json:"threadId"`
	TurnID   string        `json:"turnId"`
	Item     appServerItem `json:"item"`
}

type appServerReasoningTextDeltaNotification struct {
	ThreadID     string `json:"threadId"`
	TurnID       string `json:"turnId"`
	ItemID       string `json:"itemId"`
	Delta        string `json:"delta"`
	SummaryIndex *int   `json:"summaryIndex,omitempty"`
	ContentIndex *int   `json:"contentIndex,omitempty"`
}

type appServerItemDeltaNotification struct {
	ThreadID string `json:"threadId"`
	TurnID   string `json:"turnId"`
	ItemID   string `json:"itemId"`
	Delta    string `json:"delta"`
}

type appServerErrorNotification struct {
	ThreadID  string          `json:"threadId"`
	TurnID    string          `json:"turnId"`
	WillRetry bool            `json:"willRetry"`
	Error     json.RawMessage `json:"error"`
}

type appServerItem struct {
	Type              string                                      `json:"type"`
	ID                string                                      `json:"id"`
	Status            string                                      `json:"status,omitempty"`
	Kind              string                                      `json:"kind,omitempty"`
	ClientID          string                                      `json:"clientId,omitempty"`
	CallID            string                                      `json:"callId,omitempty"`
	CallIDSnake       string                                      `json:"call_id,omitempty"`
	Text              string                                      `json:"text,omitempty"`
	Content           json.RawMessage                             `json:"content,omitempty"`
	Summary           json.RawMessage                             `json:"summary,omitempty"`
	Output            json.RawMessage                             `json:"output,omitempty"`
	Result            json.RawMessage                             `json:"result,omitempty"`
	ContentItems      []appServerDynamicToolCallOutputContentItem `json:"contentItems,omitempty"`
	Server            string                                      `json:"server,omitempty"`
	Namespace         string                                      `json:"namespace,omitempty"`
	Tool              string                                      `json:"tool,omitempty"`
	SenderThreadID    string                                      `json:"senderThreadId,omitempty"`
	ReceiverThreadIDs []string                                    `json:"receiverThreadIds,omitempty"`
	AgentThreadID     string                                      `json:"agentThreadId,omitempty"`
	AgentPath         string                                      `json:"agentPath,omitempty"`
	Prompt            string                                      `json:"prompt,omitempty"`
	Model             string                                      `json:"model,omitempty"`
	ReasoningEffort   string                                      `json:"reasoningEffort,omitempty"`
	AgentsStates      map[string]appServerCollabAgentState        `json:"agentsStates,omitempty"`
	Arguments         json.RawMessage                             `json:"arguments,omitempty"`
	Action            json.RawMessage                             `json:"action,omitempty"`
	Query             string                                      `json:"query,omitempty"`
	Command           string                                      `json:"command,omitempty"`
	Changes           []appServerChange                           `json:"changes,omitempty"`
}

type appServerCollabAgentState struct {
	Status  string `json:"status"`
	Message string `json:"message,omitempty"`
}

type appServerChange struct {
	Path string              `json:"path"`
	Kind appServerChangeKind `json:"kind"`
}

type appServerChangeKind string

func (kind *appServerChangeKind) UnmarshalJSON(data []byte) error {
	var legacy string
	if err := json.Unmarshal(data, &legacy); err == nil {
		*kind = appServerChangeKind(strings.TrimSpace(legacy))
		return nil
	}

	var tagged struct {
		Type string `json:"type"`
	}
	if err := json.Unmarshal(data, &tagged); err != nil {
		return fmt.Errorf("decode codex file change kind: %w", err)
	}
	if strings.TrimSpace(tagged.Type) == "" {
		return fmt.Errorf("decode codex file change kind: missing type")
	}
	*kind = appServerChangeKind(strings.TrimSpace(tagged.Type))
	return nil
}

// ---------------------------------------------------------------------------
// Translated event (internal pipe format)
// ---------------------------------------------------------------------------

type translatedEvent struct {
	Kind           string         `json:"kind"`
	Text           string         `json:"text,omitempty"`
	ToolName       string         `json:"toolName,omitempty"`
	ToolCallID     string         `json:"toolCallId,omitempty"`
	ToolInput      interface{}    `json:"toolInput,omitempty"`
	PromptKind     string         `json:"promptKind,omitempty"`
	Runtime        string         `json:"runtime,omitempty"`
	EventType      string         `json:"eventType,omitempty"`
	SessionID      string         `json:"sessionId,omitempty"`
	RuntimeInputID string         `json:"runtimeInputId,omitempty"`
	Message        string         `json:"message,omitempty"`
	Title          string         `json:"title,omitempty"`
	SectionID      string         `json:"sectionId,omitempty"`
	Metadata       map[string]any `json:"metadata,omitempty"`
	// StopReason mirrors the drivercontract.ParsedEvent.StopReason field for struct parity.
	// codex-app-server uses its own drivercontract.IdleCompactionDriver for context recovery
	// and does not consume this field via the drivercontract.ContextWindowExceededDriver path.
	StopReason              string `json:"stopReason,omitempty"`
	InputTokens             uint64 `json:"inputTokens,omitempty"`
	CachedInputTokens       uint64 `json:"cachedInputTokens,omitempty"`
	OutputTokens            uint64 `json:"outputTokens,omitempty"`
	ActiveContextTokens     uint64 `json:"activeContextTokens,omitempty"`
	ContextWindowTokens     uint64 `json:"contextWindowTokens,omitempty"`
	Model                   string `json:"model,omitempty"`
	ModelInvocationID       string `json:"modelInvocationId,omitempty"`
	ContextExposureKnown    bool   `json:"contextExposureKnown,omitempty"`
	ContextExposureExcluded bool   `json:"contextExposureExcluded,omitempty"`
	ContextCompacted        bool   `json:"contextCompacted,omitempty"`
}