package agentstate

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

const ProtocolVersion = 2

type ExecutionPhase string

const (
	ExecutionIdle       ExecutionPhase = "idle"
	ExecutionRunning    ExecutionPhase = "running"
	ExecutionRetryWait  ExecutionPhase = "retry_wait"
	ExecutionRecovering ExecutionPhase = "recovering"
	ExecutionFailed     ExecutionPhase = "failed"
)

type RunningKind string

const (
	RunningThinking RunningKind = "thinking"
	RunningWorking  RunningKind = "working"
)

type ExecutionReasonCode string

const (
	ExecutionReasonRateLimited             ExecutionReasonCode = "rate_limited"
	ExecutionReasonUpstreamUnavailable     ExecutionReasonCode = "upstream_unavailable"
	ExecutionReasonProcessInterrupted      ExecutionReasonCode = "process_interrupted"
	ExecutionReasonProgressTimeout         ExecutionReasonCode = "progress_timeout"
	ExecutionReasonRuntimeFailed           ExecutionReasonCode = "runtime_failed"
	ExecutionReasonRecoveryBudgetExhausted ExecutionReasonCode = "recovery_budget_exhausted"
)

type ActionCode string

const (
	ActionReadFile          ActionCode = "read_file"
	ActionModifyFile        ActionCode = "modify_file"
	ActionSendMessage       ActionCode = "send_message"
	ActionSendMessageOutput ActionCode = "send_message_output"
	ActionRunCommand        ActionCode = "run_command"
	ActionCommandOutput     ActionCode = "command_output"
	ActionSearch            ActionCode = "search"
	ActionInspectResource   ActionCode = "inspect_resource"
)

type AgentExecutionState struct {
	Phase        ExecutionPhase      `json:"phase"`
	RunningKind  RunningKind         `json:"runningKind,omitempty"`
	ActionCode   ActionCode          `json:"actionCode,omitempty"`
	ReasonCode   ExecutionReasonCode `json:"reasonCode,omitempty"`
	AttemptToken string              `json:"attemptToken,omitempty"`
	ChangedAt    time.Time           `json:"changedAt"`
}

func (state AgentExecutionState) Validate() error {
	if state.ChangedAt.IsZero() {
		return errors.New("execution changed time is required")
	}
	switch state.Phase {
	case ExecutionIdle:
		if state.RunningKind != "" || state.ActionCode != "" || state.ReasonCode != "" {
			return errors.New("idle execution carries active or failure fields")
		}
	case ExecutionRunning:
		if state.RunningKind != RunningThinking && state.RunningKind != RunningWorking {
			return fmt.Errorf("running execution has invalid running kind %q", state.RunningKind)
		}
		if state.ReasonCode != "" {
			return errors.New("running execution carries a reason code")
		}
		if state.ActionCode != "" && !validActionCode(state.ActionCode) {
			return fmt.Errorf("running execution has invalid action code %q", state.ActionCode)
		}
	case ExecutionRetryWait, ExecutionRecovering, ExecutionFailed:
		if !validExecutionReason(state.ReasonCode) {
			return fmt.Errorf("%s execution has invalid reason code %q", state.Phase, state.ReasonCode)
		}
		if state.RunningKind != "" || state.ActionCode != "" {
			return fmt.Errorf("%s execution carries running fields", state.Phase)
		}
	default:
		return fmt.Errorf("invalid execution phase %q", state.Phase)
	}
	return nil
}

func validExecutionReason(reason ExecutionReasonCode) bool {
	switch reason {
	case ExecutionReasonRateLimited,
		ExecutionReasonUpstreamUnavailable,
		ExecutionReasonProcessInterrupted,
		ExecutionReasonProgressTimeout,
		ExecutionReasonRuntimeFailed,
		ExecutionReasonRecoveryBudgetExhausted:
		return true
	default:
		return false
	}
}

func validActionCode(action ActionCode) bool {
	switch action {
	case ActionReadFile,
		ActionModifyFile,
		ActionSendMessage,
		ActionSendMessageOutput,
		ActionRunCommand,
		ActionCommandOutput,
		ActionSearch,
		ActionInspectResource:
		return true
	default:
		return false
	}
}

type WakeDeliveryPhase string

const (
	WakePhaseReceived       WakeDeliveryPhase = "received"
	WakePhaseReplyConfirmed WakeDeliveryPhase = "reply_confirmed"
	WakePhaseNoReply        WakeDeliveryPhase = "no_reply"
	WakePhaseFailed         WakeDeliveryPhase = "failed"
)

type WakeFailureCode string

const (
	WakeFailurePlainTextNotDelivered         WakeFailureCode = "plain_text_not_delivered"
	WakeFailureBusinessOperationRejected     WakeFailureCode = "business_operation_rejected"
	WakeFailureBusinessOperationNotCommitted WakeFailureCode = "business_operation_not_committed"
	WakeFailurePartialDeliveryFailed         WakeFailureCode = "partial_delivery_failed"
	WakeFailureRuntimeFailed                 WakeFailureCode = "runtime_failed"
	WakeFailureRecoveryBudgetExhausted       WakeFailureCode = "recovery_budget_exhausted"
)

type WakeNoReplyCode string

const (
	WakeNoReplyDecided WakeNoReplyCode = "no_reply_decided"
)

type WakeDeliveryState struct {
	LocalWakeID       string
	SourceMessageID   string
	TargetKey         string
	AckedAt           time.Time
	ChangedAt         time.Time
	Phase             WakeDeliveryPhase
	ReplyObjectIDs    []string
	NoReplyReasonCode WakeNoReplyCode
	FailureReasonCode WakeFailureCode
}

func (state WakeDeliveryState) Validate() error {
	if strings.TrimSpace(state.LocalWakeID) == "" {
		return errors.New("local wake id is required")
	}
	if strings.TrimSpace(state.SourceMessageID) == "" {
		return errors.New("source message id is required")
	}
	if strings.TrimSpace(state.TargetKey) == "" {
		return errors.New("target key is required")
	}
	if state.AckedAt.IsZero() {
		return errors.New("acked time is required")
	}
	if state.ChangedAt.IsZero() {
		return errors.New("wake changed time is required")
	}
	if err := validateReplyObjectIDs(state.ReplyObjectIDs); err != nil {
		return err
	}
	switch state.Phase {
	case WakePhaseReceived:
		if len(state.ReplyObjectIDs) > 0 || state.NoReplyReasonCode != "" || state.FailureReasonCode != "" {
			return errors.New("received wake carries terminal evidence")
		}
	case WakePhaseReplyConfirmed:
		if len(state.ReplyObjectIDs) == 0 {
			return errors.New("reply_confirmed wake requires reply object ids")
		}
		if state.NoReplyReasonCode != "" || state.FailureReasonCode != "" {
			return errors.New("reply_confirmed wake carries a reason code")
		}
	case WakePhaseNoReply:
		if !validWakeNoReplyReason(state.NoReplyReasonCode) {
			return fmt.Errorf("no_reply wake has invalid reason code %q", state.NoReplyReasonCode)
		}
		if len(state.ReplyObjectIDs) > 0 || state.FailureReasonCode != "" {
			return errors.New("no_reply wake carries conflicting terminal evidence")
		}
	case WakePhaseFailed:
		if !validWakeFailureReason(state.FailureReasonCode) {
			return fmt.Errorf("failed wake has invalid reason code %q", state.FailureReasonCode)
		}
		if len(state.ReplyObjectIDs) > 0 || state.NoReplyReasonCode != "" {
			return errors.New("failed wake carries conflicting terminal evidence")
		}
	default:
		return fmt.Errorf("invalid wake delivery phase %q", state.Phase)
	}
	return nil
}

func validateReplyObjectIDs(ids []string) error {
	seen := make(map[string]struct{}, len(ids))
	for _, id := range ids {
		trimmed := strings.TrimSpace(id)
		if trimmed == "" {
			return errors.New("reply object id is empty")
		}
		if _, ok := seen[trimmed]; ok {
			return fmt.Errorf("duplicate reply object id %q", trimmed)
		}
		seen[trimmed] = struct{}{}
	}
	return nil
}

func validWakeNoReplyReason(reason WakeNoReplyCode) bool {
	return reason == WakeNoReplyDecided
}

func validWakeFailureReason(reason WakeFailureCode) bool {
	switch reason {
	case WakeFailurePlainTextNotDelivered,
		WakeFailureBusinessOperationRejected,
		WakeFailureBusinessOperationNotCommitted,
		WakeFailurePartialDeliveryFailed,
		WakeFailureRuntimeFailed,
		WakeFailureRecoveryBudgetExhausted:
		return true
	default:
		return false
	}
}

type wakeDeliveryJSON struct {
	LocalWakeID     string            `json:"localWakeId"`
	SourceMessageID string            `json:"sourceMessageId"`
	TargetKey       string            `json:"targetKey"`
	AckedAt         time.Time         `json:"surfacedAt"`
	ChangedAt       time.Time         `json:"changedAt"`
	Phase           WakeDeliveryPhase `json:"phase"`
	ReplyObjectIDs  []string          `json:"replyObjectIds,omitempty"`
	ReasonCode      string            `json:"reasonCode,omitempty"`
}

func (state WakeDeliveryState) MarshalJSON() ([]byte, error) {
	if err := state.Validate(); err != nil {
		return nil, err
	}
	reason := ""
	switch state.Phase {
	case WakePhaseNoReply:
		reason = string(state.NoReplyReasonCode)
	case WakePhaseFailed:
		reason = string(state.FailureReasonCode)
	}
	return json.Marshal(wakeDeliveryJSON{
		LocalWakeID:     state.LocalWakeID,
		SourceMessageID: state.SourceMessageID,
		TargetKey:       state.TargetKey,
		AckedAt:         state.AckedAt,
		ChangedAt:       state.ChangedAt,
		Phase:           state.Phase,
		ReplyObjectIDs:  state.ReplyObjectIDs,
		ReasonCode:      reason,
	})
}

func (state *WakeDeliveryState) UnmarshalJSON(data []byte) error {
	var wire wakeDeliveryJSON
	if err := json.Unmarshal(data, &wire); err != nil {
		return err
	}
	var alternate struct {
		AckedAt time.Time `json:"ackedAt"`
	}
	_ = json.Unmarshal(data, &alternate)
	if wire.AckedAt.IsZero() {
		wire.AckedAt = alternate.AckedAt
	}
	restored := WakeDeliveryState{
		LocalWakeID:     wire.LocalWakeID,
		SourceMessageID: wire.SourceMessageID,
		TargetKey:       wire.TargetKey,
		AckedAt:         wire.AckedAt,
		ChangedAt:       wire.ChangedAt,
		Phase:           wire.Phase,
		ReplyObjectIDs:  append([]string(nil), wire.ReplyObjectIDs...),
	}
	switch wire.Phase {
	case WakePhaseNoReply:
		restored.NoReplyReasonCode = WakeNoReplyCode(wire.ReasonCode)
	case WakePhaseFailed:
		restored.FailureReasonCode = WakeFailureCode(wire.ReasonCode)
	default:
		if wire.ReasonCode != "" {
			return fmt.Errorf("%s wake carries reason code %q", wire.Phase, wire.ReasonCode)
		}
	}
	if err := restored.Validate(); err != nil {
		return err
	}
	*state = restored
	return nil
}

type DiagnosticCategory string

const (
	DiagnosticExecution DiagnosticCategory = "execution"
	DiagnosticDelivery  DiagnosticCategory = "delivery"
)

type DeliveryDiagnosticReason string

const DeliveryDiagnosticReceiptReconcilePending DeliveryDiagnosticReason = "receipt_reconcile_pending"

type AgentDiagnostic struct {
	DiagnosticID      string
	TargetKey         string
	LocalWakeIDs      []string
	ClientOperationID string
	CreatedAt         time.Time
	Category          DiagnosticCategory
	Terminal          bool
	ExecutionReason   ExecutionReasonCode
	DeliveryReason    DeliveryDiagnosticReason
	FailureReason     WakeFailureCode
}

func (diagnostic AgentDiagnostic) Validate() error {
	if strings.TrimSpace(diagnostic.DiagnosticID) == "" {
		return errors.New("diagnostic id is required")
	}
	if diagnostic.CreatedAt.IsZero() {
		return errors.New("diagnostic created time is required")
	}
	if err := validateNonEmptyUniqueIDs(diagnostic.LocalWakeIDs, "local wake id"); err != nil {
		return err
	}
	switch diagnostic.Category {
	case DiagnosticExecution:
		if !validExecutionReason(diagnostic.ExecutionReason) {
			return fmt.Errorf("execution diagnostic has invalid reason code %q", diagnostic.ExecutionReason)
		}
		if diagnostic.DeliveryReason != "" || diagnostic.FailureReason != "" {
			return errors.New("execution diagnostic carries delivery reason")
		}
		if len(diagnostic.LocalWakeIDs) > 0 && strings.TrimSpace(diagnostic.TargetKey) == "" {
			return errors.New("scoped execution diagnostic requires target")
		}
	case DiagnosticDelivery:
		if strings.TrimSpace(diagnostic.TargetKey) == "" {
			return errors.New("delivery diagnostic requires target")
		}
		if diagnostic.ExecutionReason != "" {
			return errors.New("delivery diagnostic carries execution reason")
		}
		switch {
		case diagnostic.DeliveryReason == DeliveryDiagnosticReceiptReconcilePending:
			if diagnostic.Terminal {
				return errors.New("receipt reconcile diagnostic must be non-terminal")
			}
			if diagnostic.FailureReason != "" {
				return errors.New("receipt reconcile diagnostic carries failure reason")
			}
			if strings.TrimSpace(diagnostic.ClientOperationID) == "" {
				return errors.New("receipt reconcile diagnostic requires client operation id")
			}
		case validWakeFailureReason(diagnostic.FailureReason):
			if !diagnostic.Terminal {
				return errors.New("delivery failure diagnostic must be terminal")
			}
			if diagnostic.DeliveryReason != "" {
				return errors.New("delivery failure diagnostic carries reconcile reason")
			}
		default:
			return errors.New("delivery diagnostic requires a typed reason")
		}
	default:
		return fmt.Errorf("invalid diagnostic category %q", diagnostic.Category)
	}
	return nil
}

func validateNonEmptyUniqueIDs(ids []string, label string) error {
	seen := make(map[string]struct{}, len(ids))
	for _, id := range ids {
		trimmed := strings.TrimSpace(id)
		if trimmed == "" {
			return fmt.Errorf("%s is empty", label)
		}
		if _, ok := seen[trimmed]; ok {
			return fmt.Errorf("duplicate %s %q", label, trimmed)
		}
		seen[trimmed] = struct{}{}
	}
	return nil
}

type agentDiagnosticJSON struct {
	DiagnosticID      string             `json:"diagnosticId"`
	TargetKey         string             `json:"targetKey,omitempty"`
	LocalWakeIDs      []string           `json:"localWakeIds,omitempty"`
	ClientOperationID string             `json:"clientOperationId,omitempty"`
	CreatedAt         time.Time          `json:"createdAt"`
	Category          DiagnosticCategory `json:"category"`
	Terminal          bool               `json:"terminal"`
	ReasonCode        string             `json:"reasonCode"`
}

func (diagnostic AgentDiagnostic) MarshalJSON() ([]byte, error) {
	if err := diagnostic.Validate(); err != nil {
		return nil, err
	}
	reason := ""
	switch {
	case diagnostic.Category == DiagnosticExecution:
		reason = string(diagnostic.ExecutionReason)
	case diagnostic.DeliveryReason != "":
		reason = string(diagnostic.DeliveryReason)
	default:
		reason = string(diagnostic.FailureReason)
	}
	return json.Marshal(agentDiagnosticJSON{
		DiagnosticID:      diagnostic.DiagnosticID,
		TargetKey:         diagnostic.TargetKey,
		LocalWakeIDs:      diagnostic.LocalWakeIDs,
		ClientOperationID: diagnostic.ClientOperationID,
		CreatedAt:         diagnostic.CreatedAt,
		Category:          diagnostic.Category,
		Terminal:          diagnostic.Terminal,
		ReasonCode:        reason,
	})
}

func (diagnostic *AgentDiagnostic) UnmarshalJSON(data []byte) error {
	var wire agentDiagnosticJSON
	if err := json.Unmarshal(data, &wire); err != nil {
		return err
	}
	restored := AgentDiagnostic{
		DiagnosticID:      wire.DiagnosticID,
		TargetKey:         wire.TargetKey,
		LocalWakeIDs:      append([]string(nil), wire.LocalWakeIDs...),
		ClientOperationID: wire.ClientOperationID,
		CreatedAt:         wire.CreatedAt,
		Category:          wire.Category,
		Terminal:          wire.Terminal,
	}
	switch wire.Category {
	case DiagnosticExecution:
		restored.ExecutionReason = ExecutionReasonCode(wire.ReasonCode)
	case DiagnosticDelivery:
		if wire.ReasonCode == string(DeliveryDiagnosticReceiptReconcilePending) {
			restored.DeliveryReason = DeliveryDiagnosticReceiptReconcilePending
		} else {
			restored.FailureReason = WakeFailureCode(wire.ReasonCode)
		}
	}
	if err := restored.Validate(); err != nil {
		return err
	}
	*diagnostic = restored
	return nil
}

type EnvelopeMode string

const (
	SnapshotMode EnvelopeMode = "snapshot"
	DeltaMode    EnvelopeMode = "delta"
)

type EnvelopeBaseV2 struct {
	Version      int       `json:"version"`
	AgentID      string    `json:"agentId"`
	SessionEpoch uint64    `json:"sessionEpoch"`
	Revision     uint64    `json:"revision"`
	EmittedAt    time.Time `json:"emittedAt"`
}

func (base EnvelopeBaseV2) validate() error {
	if base.Version != ProtocolVersion {
		return fmt.Errorf("invalid agent state protocol version %d", base.Version)
	}
	if strings.TrimSpace(base.AgentID) == "" {
		return errors.New("agent id is required")
	}
	if base.SessionEpoch == 0 {
		return errors.New("session epoch is required")
	}
	if base.Revision == 0 {
		return errors.New("revision is required")
	}
	if base.EmittedAt.IsZero() {
		return errors.New("emitted time is required")
	}
	return nil
}

type AgentStateSnapshotV2 struct {
	EnvelopeBaseV2
	Mode        EnvelopeMode        `json:"mode"`
	Execution   AgentExecutionState `json:"execution"`
	Deliveries  []WakeDeliveryState `json:"deliveries"`
	Diagnostics []AgentDiagnostic   `json:"diagnostics"`
}

func (snapshot AgentStateSnapshotV2) Validate() error {
	if err := snapshot.EnvelopeBaseV2.validate(); err != nil {
		return err
	}
	if snapshot.Mode != SnapshotMode {
		return fmt.Errorf("snapshot has invalid mode %q", snapshot.Mode)
	}
	if err := snapshot.Execution.Validate(); err != nil {
		return fmt.Errorf("validate snapshot execution: %w", err)
	}
	if err := validateDeliveries(snapshot.Deliveries); err != nil {
		return err
	}
	if err := validateDiagnostics(snapshot.Diagnostics); err != nil {
		return err
	}
	return nil
}

type AgentStateDeltaV2 struct {
	EnvelopeBaseV2
	Mode               EnvelopeMode         `json:"mode"`
	BaseRevision       uint64               `json:"baseRevision"`
	Execution          *AgentExecutionState `json:"execution,omitempty"`
	DeliveryUpserts    []WakeDeliveryState  `json:"deliveryUpserts,omitempty"`
	DeliveryRemovals   []string             `json:"deliveryRemovals,omitempty"`
	DiagnosticUpserts  []AgentDiagnostic    `json:"diagnosticUpserts,omitempty"`
	DiagnosticRemovals []string             `json:"diagnosticRemovals,omitempty"`
}

func (delta AgentStateDeltaV2) Validate() error {
	if err := delta.EnvelopeBaseV2.validate(); err != nil {
		return err
	}
	if delta.Mode != DeltaMode {
		return fmt.Errorf("delta has invalid mode %q", delta.Mode)
	}
	if delta.BaseRevision == 0 || delta.BaseRevision >= delta.Revision {
		return fmt.Errorf("delta has invalid base revision %d", delta.BaseRevision)
	}
	if delta.Execution != nil {
		if err := delta.Execution.Validate(); err != nil {
			return fmt.Errorf("validate delta execution: %w", err)
		}
	}
	if err := validateDeliveries(delta.DeliveryUpserts); err != nil {
		return err
	}
	if err := validateNonEmptyUniqueIDs(delta.DeliveryRemovals, "delivery removal"); err != nil {
		return err
	}
	if err := validateDiagnostics(delta.DiagnosticUpserts); err != nil {
		return err
	}
	if err := validateNonEmptyUniqueIDs(delta.DiagnosticRemovals, "diagnostic removal"); err != nil {
		return err
	}
	if delta.Execution == nil &&
		len(delta.DeliveryUpserts) == 0 &&
		len(delta.DeliveryRemovals) == 0 &&
		len(delta.DiagnosticUpserts) == 0 &&
		len(delta.DiagnosticRemovals) == 0 {
		return errors.New("delta has no changes")
	}
	return nil
}

func validateDeliveries(deliveries []WakeDeliveryState) error {
	seen := make(map[string]struct{}, len(deliveries))
	for i, delivery := range deliveries {
		if err := delivery.Validate(); err != nil {
			return fmt.Errorf("validate delivery %d: %w", i, err)
		}
		if _, ok := seen[delivery.LocalWakeID]; ok {
			return fmt.Errorf("duplicate delivery %q", delivery.LocalWakeID)
		}
		seen[delivery.LocalWakeID] = struct{}{}
	}
	return nil
}

func validateDiagnostics(diagnostics []AgentDiagnostic) error {
	seen := make(map[string]struct{}, len(diagnostics))
	for i, diagnostic := range diagnostics {
		if err := diagnostic.Validate(); err != nil {
			return fmt.Errorf("validate diagnostic %d: %w", i, err)
		}
		if _, ok := seen[diagnostic.DiagnosticID]; ok {
			return fmt.Errorf("duplicate diagnostic %q", diagnostic.DiagnosticID)
		}
		seen[diagnostic.DiagnosticID] = struct{}{}
	}
	return nil
}

type State struct {
	Execution   AgentExecutionState
	Deliveries  []WakeDeliveryState
	Diagnostics []AgentDiagnostic
}