// Package protocol implements the correx WebSocket wire contract used to talk
// to apps/server over /stream. The server (Kotlin/kotlinx.serialization) uses a
// "type" class discriminator. Server->client variants carry @SerialName values
// (e.g. "session.announced"); client->server variants have no @SerialName, so the
// discriminator is the fully-qualified Kotlin class name.
package protocol

import "encoding/json"

// clientPrefix is the kotlinx default serial name for nested ClientMessage
// subclasses (qualified class name). The server decodes against this exact value.
const clientPrefix = "com.correx.apps.server.protocol.ClientMessage."

// Server message discriminator values (@SerialName on ServerMessage variants).
const (
	TypeSessionAnnounced = "session.announced"
	TypeChatTurn         = "chat.turn"
	TypeSessionPaused    = "session.paused"
	TypeSessionResumed   = "session.resumed"
	TypeSessionCompleted = "session.completed"
	TypeSessionFailed    = "session.failed"
	TypeSessionSnapshot  = "session_snapshot"
	TypeStageStarted     = "stage.started"
	TypeStageCompleted   = "stage.completed"
	TypeStageFailed      = "stage.failed"
	TypeInferenceStarted = "inference.started"
	TypeInferenceDone    = "inference.completed"
	TypeInferenceTimeout = "inference.timed_out"
	TypeInferenceFailed  = "inference.failed"
	TypeInferenceRetry   = "inference.retry"
	TypeToolStarted      = "tool.started"
	TypeToolCompleted    = "tool.completed"
	TypeToolFailed       = "tool.failed"
	TypeToolRejected     = "tool.rejected"
	TypeApprovalRequired = "approval.required"
	TypeApprovalResolved = "approval.resolved"
	TypeClarification    = "clarification.required"
	TypeWorkflowProposed = "workflow.proposed"
	TypeWorkspaceBound   = "session.workspace_bound"
	TypeStageManifest    = "stage.tool_manifest"
	TypeProviderStatus   = "provider.status_changed"
	TypeProtocolError    = "protocol_error"
	TypeArtifactCreated  = "artifact.created"
	TypeArtifactValid    = "artifact.validated"
	TypePlanLocked       = "plan.locked"
	TypeRouterResponse   = "router.response"
	TypeSnapshotComplete = "snapshot_complete"
	TypeWorkflowList     = "workflow.list"
	TypeToolAssessed     = "tool.assessed"
	TypeModelChanged     = "model.changed"
	TypeModelList        = "model.list"
	TypeResourceStatus   = "resource.status"
	TypeRouterNarration  = "router.narration"
	TypeArtifactList     = "artifact.list"
	TypeConfigSnapshot   = "config.snapshot"
	TypeSessionStats     = "session.stats"
	TypeIdeaList         = "idea.list"
)

// ServerMessage is a flat decode of every server->client variant. Field names
// across variants don't collide in incompatible JSON types, so a single struct
// keyed by Type decodes the whole protocol. Value-class IDs serialize as plain
// strings (verified against ServerMessageSerializationTest).
type ServerMessage struct {
	Type string `json:"type"`

	SessionID     string  `json:"sessionId"`
	WorkflowID    string  `json:"workflowId"`
	StageID       string  `json:"stageId"`
	ToolName      string  `json:"toolName"`
	Role          string  `json:"role"`
	TurnID        string  `json:"turnId"`
	Reason        string  `json:"reason"`
	Summary       string  `json:"outputSummary"`
	Response      string  `json:"responseText"`
	LastOutput    string  `json:"lastOutput"`   // session_snapshot: last stage output (reopen)
	LastResponse  string  `json:"lastResponse"` // session_snapshot: last stage response (reopen)
	Content       string  `json:"content"`
	Diff          *string `json:"diff"`
	Tier          string  `json:"tier"`
	Message       string  `json:"message"`
	RequestID     string  `json:"requestId"`
	ArtifactID    string  `json:"artifactId"`
	ProviderID    string  `json:"providerId"`
	Preview       *string `json:"preview"`
	Disposition   string  `json:"disposition"`
	Outcome       string  `json:"outcome"`       // approval.resolved: APPROVED | REJECTED | AUTO_APPROVED
	WorkspaceRoot string  `json:"workspaceRoot"` // session.workspace_bound: the bound cwd

	// plan.locked: the compiled phase-2 stage ids in execution order
	StageIDs []string `json:"stageIds"`

	// inference.retry
	AttemptNumber int    `json:"attemptNumber"`
	MaxAttempts   int    `json:"maxAttempts"`
	FailureReason string `json:"failureReason"`

	SteeringEmitted bool  `json:"steeringEmitted"`
	OccurredAt      int64 `json:"occurredAt"`
	ElapsedMs       int64 `json:"elapsedMs"`
	LastSequence    int64 `json:"lastSequence"`

	// model.changed
	ModelID string `json:"modelId"`
	Loaded  bool   `json:"loaded"`

	// model.list (current reuses the field below)
	Models  []string `json:"models"`
	Current string   `json:"current"`

	// chat.turn router metrics (nil when not a ROUTER turn or server omits them)
	LatencyMs   *int64 `json:"latencyMs"`
	TotalTokens *int   `json:"totalTokens"`

	// resource.status (nil = unavailable)
	GpuMemoryUsedMb   *int64 `json:"gpuMemoryUsedMb"`
	GpuMemoryTotalMb  *int64 `json:"gpuMemoryTotalMb"`
	GpuUtilizationPct *int   `json:"gpuUtilizationPct"`
	ProcessRssMb      *int64 `json:"processRssMb"`
	SystemRamUsedMb   *int64 `json:"systemRamUsedMb"`
	SystemRamTotalMb  *int64 `json:"systemRamTotalMb"`

	State          *SessionStateDto   `json:"state"`
	RiskSummary    *RiskSummaryDto    `json:"riskSummary"`
	Status         *ProviderHealthDto `json:"status"`
	PendingAppr    []ApprovalDto      `json:"pendingApprovals"`
	Tools          []ToolRecordDto    `json:"tools"`
	RecentEvents   []EventEntryDto    `json:"recentEvents"`
	Stages         []StageToolDecl    `json:"stages"`
	Workflows      []WorkflowDto      `json:"workflows"`
	AssessedIssues []AssessedIssueDto `json:"issues"`
	Artifacts      []ArtifactDto      `json:"artifacts"`

	// config.snapshot
	ConfigFields          []ConfigFieldDto `json:"fields"`
	ConfigRestartRequired []string         `json:"restartRequired"`
	ConfigError           *string          `json:"error"`

	// session.stats — derived metrics for a session (reply to GetSessionStats)
	Stats *StatsDto `json:"stats"`

	// clarification.required — open questions a stage raised for the operator
	Questions []ClarificationQuestionDto `json:"questions"`

	// workflow.proposed — the router's triage suggestion of candidate workflows
	ProposalID      string                `json:"proposalId"`
	Prompt          string                `json:"prompt"`
	OriginalRequest string                `json:"originalRequest"`
	Candidates      []ProposedWorkflowDto `json:"candidates"`

	// idea.list — the cross-session idea board
	Ideas []IdeaDto `json:"ideas"`
}

// IdeaDto is one idea on the cross-session board. Mirrors the Kotlin IdeaDto.
type IdeaDto struct {
	IdeaID       string `json:"ideaId"`
	Text         string `json:"text"`
	CapturedAtMs int64  `json:"capturedAtMs"`
}

// ProposedWorkflowDto is one candidate workflow the router suggests, with a short
// reason. Mirrors the Kotlin ProposedWorkflow (core:events), the wire shape.
type ProposedWorkflowDto struct {
	WorkflowID string `json:"workflowId"`
	Reason     string `json:"reason"`
}

// ClarificationQuestionDto is one open question from a stage. Empty Options ->
// free-text answer; non-empty -> selectable choices. Mirrors the Kotlin
// ClarificationQuestion (core:events), which is the wire shape.
type ClarificationQuestionDto struct {
	ID          string   `json:"id"`
	Prompt      string   `json:"prompt"`
	Options     []string `json:"options"`
	MultiSelect bool     `json:"multiSelect"`
	Header      string   `json:"header"`
}

// ClarificationAnswerDto is the operator's answer to one question.
type ClarificationAnswerDto struct {
	QuestionID string `json:"questionId"`
	Value      string `json:"value"`
}

// StatsDto mirrors the server's MetricsReport: a session's derived metrics
// (observability-spec §2/§3). Raw sums plus the derived ratios the pane renders.
type StatsDto struct {
	SessionID          string                 `json:"sessionId"`
	EventCount         int64                  `json:"eventCount"`
	SessionDurationMs  int64                  `json:"sessionDurationMs"`
	InferenceCount     int64                  `json:"inferenceCount"`
	InferenceMs        int64                  `json:"inferenceMs"`
	PromptTokens       int64                  `json:"promptTokens"`
	CompletionTokens   int64                  `json:"completionTokens"`
	TokensPerSecond    float64                `json:"tokensPerSecond"`
	PerProvider        []ProviderStatsDto     `json:"perProvider"`
	ToolCount          int64                  `json:"toolCount"`
	ToolMs             int64                  `json:"toolMs"`
	PerTool            []ToolStatsDto         `json:"perTool"`
	ApprovalsRequested int64                  `json:"approvalsRequested"`
	ApprovalsResolved  int64                  `json:"approvalsResolved"`
	ApprovalsPending   int64                  `json:"approvalsPending"`
	ApprovalWaitMs     int64                  `json:"approvalWaitMs"`
	AvgApprovalWaitMs  int64                  `json:"avgApprovalWaitMs"`
	PerTier            []TierApprovalStatsDto `json:"perTier"`
	Failures           FailureMetricsDto      `json:"failures"`
	InferencePct       float64                `json:"inferencePct"`
	ToolPct            float64                `json:"toolPct"`
	ApprovalWaitPct    float64                `json:"approvalWaitPct"`
}

type ProviderStatsDto struct {
	Provider         string  `json:"provider"`
	CompletedCount   int64   `json:"completedCount"`
	TotalLatencyMs   int64   `json:"totalLatencyMs"`
	PromptTokens     int64   `json:"promptTokens"`
	CompletionTokens int64   `json:"completionTokens"`
	TokensPerSecond  float64 `json:"tokensPerSecond"`
}

type ToolStatsDto struct {
	ToolName        string `json:"toolName"`
	CompletedCount  int64  `json:"completedCount"`
	TotalDurationMs int64  `json:"totalDurationMs"`
}

type TierApprovalStatsDto struct {
	Tier           string `json:"tier"`
	RequestedCount int64  `json:"requestedCount"`
	ResolvedCount  int64  `json:"resolvedCount"`
	TotalWaitMs    int64  `json:"totalWaitMs"`
	AvgWaitMs      int64  `json:"avgWaitMs"`
}

type FailureMetricsDto struct {
	InferenceFailures int64 `json:"inferenceFailures"`
	InferenceTimeouts int64 `json:"inferenceTimeouts"`
	ToolFailures      int64 `json:"toolFailures"`
	ToolRejections    int64 `json:"toolRejections"`
	StageFailures     int64 `json:"stageFailures"`
	WorkflowFailures  int64 `json:"workflowFailures"`
}

type ConfigFieldDto struct {
	Key             string   `json:"key"`
	Type            string   `json:"type"`
	Value           string   `json:"value"`
	Options         []string `json:"options"`
	RestartRequired bool     `json:"restartRequired"`
}

type ArtifactDto struct {
	ArtifactID    string  `json:"artifactId"`
	StageID       string  `json:"stageId"`
	SchemaVersion int     `json:"schemaVersion"`
	Phase         string  `json:"phase"`
	Content       *string `json:"content"`
}

type SessionStateDto struct {
	Status         string  `json:"status"`
	CurrentStageID *string `json:"currentStageId"`
	PauseReason    *string `json:"pauseReason"`
}

type RiskSummaryDto struct {
	Level             string   `json:"level"`
	Factors           []string `json:"factors"`
	RecommendedAction string   `json:"recommendedAction"`
	Rationale         []string `json:"rationale"`
}

type AssessedIssueDto struct {
	Code     string `json:"code"`
	Message  string `json:"message"`
	Severity string `json:"severity"`
}

type ProviderHealthDto struct {
	ProviderID  string `json:"providerId"`
	Status      string `json:"status"`
	LoadPercent *int   `json:"loadPercent"`
}

type ApprovalDto struct {
	RequestID string  `json:"requestId"`
	Tier      string  `json:"tier"`
	ToolName  *string `json:"toolName"`
	Preview   *string `json:"preview"`
}

type StageToolDecl struct {
	StageID string     `json:"stageId"`
	Tools   []ToolDecl `json:"tools"`
}

type ToolDecl struct {
	Name string `json:"name"`
	Tier int    `json:"tier"`
}

type EventEntryDto struct {
	Timestamp int64  `json:"timestamp"`
	Type      string `json:"type"`
	Detail    string `json:"detail"`
}

type ToolRecordDto struct {
	Name   string  `json:"name"`
	Tier   int     `json:"tier"`
	Status string  `json:"status"`
	Diff   *string `json:"diff"`
}

type WorkflowDto struct {
	WorkflowID  string `json:"workflowId"`
	Description string `json:"description"`
}

// Decode parses one server frame into a ServerMessage.
func Decode(raw []byte) (ServerMessage, error) {
	var m ServerMessage
	err := json.Unmarshal(raw, &m)
	return m, err
}

// IsEventBearing reports whether a message carries a session event (and so must
// be buffered during the snapshot phase). Control/snapshot messages return false
// and are applied immediately.
func (m ServerMessage) IsEventBearing() bool {
	switch m.Type {
	case TypeSessionAnnounced, TypeChatTurn, TypeSessionPaused, TypeSessionResumed,
		TypeSessionCompleted, TypeSessionFailed,
		TypeStageStarted, TypeStageCompleted, TypeStageFailed,
		TypeInferenceStarted, TypeInferenceDone, TypeInferenceTimeout,
		TypeInferenceFailed, TypeInferenceRetry,
		TypeToolStarted, TypeToolCompleted, TypeToolFailed, TypeToolRejected,
		TypeToolAssessed,
		TypeApprovalRequired, TypeApprovalResolved, TypeClarification,
		TypeWorkflowProposed,
		TypeWorkspaceBound, TypeArtifactCreated,
		TypeArtifactValid, TypePlanLocked,
		TypeRouterNarration:
		return true
	default:
		return false
	}
}

// --- Client -> server encoders ---

func encode(typeName string, fields map[string]any) []byte {
	fields["type"] = clientPrefix + typeName
	b, _ := json.Marshal(fields)
	return b
}

// StartChatSession opens a new chat session with a client-generated sessionId.
func StartChatSession(sessionID, text string) []byte {
	return encode("StartChatSession", map[string]any{"sessionId": sessionID, "text": text})
}

// ChatInput sends a turn into an existing session. mode is "CHAT" or "STEERING".
func ChatInput(sessionID, text, mode string) []byte {
	return encode("ChatInput", map[string]any{"sessionId": sessionID, "text": text, "mode": mode})
}

// StartSession launches a workflow-backed session. input is the optional freeform request
// seeded into the decision journal; pass "" for fixed-task workflows (e.g. healthcheck).
func StartSession(workflowID, input string) []byte {
	msg := map[string]any{"workflowId": workflowID, "config": nil}
	if input != "" {
		msg["input"] = input
	}
	return encode("StartSession", msg)
}

// CancelSession cancels a running session.
func CancelSession(sessionID string) []byte {
	return encode("CancelSession", map[string]any{"sessionId": sessionID})
}

// ListArtifacts asks the server for the artifacts produced by a session, with content.
func ListArtifacts(sessionID string) []byte {
	return encode("ListArtifacts", map[string]any{"sessionId": sessionID})
}

// GetSessionStats asks the server for a session's derived metrics (replied to with session.stats).
func GetSessionStats(sessionID string) []byte {
	return encode("GetSessionStats", map[string]any{"sessionId": sessionID})
}

// GetConfig requests the current editable config (replied to with a config.snapshot).
func GetConfig() []byte {
	return encode("GetConfig", map[string]any{})
}

// ListIdeas requests the cross-session idea board (replied to with an idea.list).
func ListIdeas() []byte {
	return encode("ListIdeas", map[string]any{})
}

// DiscardIdea removes an idea from the board (server tombstones it; replies with a fresh idea.list).
func DiscardIdea(ideaID string) []byte {
	return encode("DiscardIdea", map[string]any{"ideaId": ideaID})
}

// UpdateConfig asks the server to apply and persist a config patch (key -> string value).
func UpdateConfig(patch map[string]string) []byte {
	return encode("UpdateConfig", map[string]any{"patch": patch})
}

// ApprovalResponse answers an approval gate. decision is "APPROVE" or "REJECT".
func ApprovalResponse(requestID, decision string, steeringNote *string) []byte {
	return encode("ApprovalResponse", map[string]any{
		"requestId":    requestID,
		"decision":     decision,
		"steeringNote": steeringNote,
	})
}

// ClarificationResponse answers a stage's open questions; the parked stage re-runs
// with the answers in context.
func ClarificationResponse(sessionID, stageID, requestID string, answers []ClarificationAnswerDto) []byte {
	return encode("ClarificationResponse", map[string]any{
		"sessionId": sessionID,
		"stageId":   stageID,
		"requestId": requestID,
		"answers":   answers,
	})
}

// Ping is a keepalive.
func Ping(ts int64) []byte {
	return encode("Ping", map[string]any{"timestamp": ts})
}

// SwapModel asks the server to make modelID resident and pin it.
func SwapModel(modelID string) []byte {
	return encode("SwapModel", map[string]any{"modelId": modelID})
}

// ClearModelPin releases the manual-swap pin so per-stage selection resumes.
func ClearModelPin() []byte {
	return encode("ClearModelPin", map[string]any{})
}

// CreateGrant pre-authorizes a tool/tier so future calls skip the approval gate.
// scope is "SESSION" or "STAGE"; permittedTiers are tier names like "T3".
func CreateGrant(sessionID, scope string, permittedTiers []string, reason, toolName string) []byte {
	return encode("CreateGrant", map[string]any{
		"sessionId":      sessionID,
		"scope":          scope,
		"stageId":        nil,
		"permittedTiers": permittedTiers,
		"reason":         reason,
		"expiresAt":      nil,
		"toolName":       toolName,
	})
}

// Hello is the first frame sent on every (re)connect. It carries the client's
// working directory so the server can bind a workspace before any session starts.
func Hello(workingDir string) []byte {
	return encode("Hello", map[string]any{"workingDir": workingDir})
}