// Package rewindcaps declares recovery capabilities for each supported runtime
// and provides encoding/decoding helpers for recovery point references.
package rewindcaps

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"
)

const (
	RuntimeClaude                 = "claude"
	RuntimeCodex                  = "codex"
	RuntimeKimi                   = "kimi"
	RuntimeOpenCode               = "opencode"
	RuntimeBuiltin                = "builtin"
	RuntimePi                     = "pi"
	StrategyCodexAppServerTurns   = "codex_app_server_turns"
	StrategyClaudeSessionResumeAt = "claude_session_resume_at"
	StrategyKimiSessionTurns      = "kimi_session_turns"
	StrategyOpenCodeSessionRevert = "opencode_session_revert"
	StrategyPiSessionTreeRewind   = "pi_session_tree_rewind"

	RecoveryBoundaryResumeAt         = "resume_at"
	RecoveryBoundaryBeforeFirstInput = "before_first_input"

	UnsupportedReasonUnknownRuntime = "unknown runtime"
)

// Capability describes a runtime's recovery support.
type Capability struct {
	Runtime                 string
	Supported               bool
	Strategy                string
	RequiresPreparedPoint   bool
	SupportsDetachedRestore bool
	SupportsLiveRestore     bool
	UnsupportedReason       string
}

// CodexAppServerRecoveryRef is the JSON recovery reference for the codex
// app-server turn rollback strategy.
type CodexAppServerRecoveryRef struct {
	ThreadID  string `json:"threadId"`
	TurnCount int    `json:"turnCount"`
}

// ClaudeSessionResumeAtRecoveryRef is the JSON recovery reference for Claude
// Code's native resume-at strategy. BoundaryKind is either resume_at, where
// BoundaryUUID must name the last transcript message before the AgentTeam input, or
// before_first_input, where the restore must create a fresh runtime session.
type ClaudeSessionResumeAtRecoveryRef struct {
	SessionID    string `json:"sessionId"`
	BoundaryUUID string `json:"boundaryUuid,omitempty"`
	BoundaryKind string `json:"boundaryKind"`
}

// OpenCodeSessionRevertRecoveryRef is the JSON recovery reference for
// OpenCode's native session revert strategy. MessageID is the OpenCode user
// message id assigned to the AgentTeam input before it enters runtime context.
type OpenCodeSessionRevertRecoveryRef struct {
	SessionID string `json:"sessionId"`
	MessageID string `json:"messageId"`
}

// KimiSessionTurnsRecoveryRef identifies the completed-turn boundary retained
// when Kimi's durable wire and context logs are restored.
type KimiSessionTurnsRecoveryRef struct {
	SessionID  string `json:"sessionId"`
	TurnCount  int    `json:"turnCount"`
	InputCount int    `json:"inputCount"`
}

// PiSessionTreeRewindRef identifies the exact session leaf entry that was
// current before the AgentTeam message being rewound. BoundaryEntryID names the
// last entry retained (inclusive); BoundaryKind distinguishes resume_at from a
// fresh-session before_first_input boundary.
type PiSessionTreeRewindRef struct {
	SessionID       string `json:"sessionId"`
	BoundaryEntryID string `json:"boundaryEntryId,omitempty"`
	BoundaryKind    string `json:"boundaryKind"`
}

var capabilityByRuntime = map[string]Capability{
	"claude": {
		Runtime:                 "claude",
		Supported:               true,
		Strategy:                StrategyClaudeSessionResumeAt,
		RequiresPreparedPoint:   true,
		SupportsDetachedRestore: true,
		SupportsLiveRestore:     false,
	},
	"codex": {
		Runtime:                 "codex",
		Supported:               true,
		Strategy:                StrategyCodexAppServerTurns,
		RequiresPreparedPoint:   true,
		SupportsDetachedRestore: true,
		SupportsLiveRestore:     true,
	},
	RuntimeKimi: {
		Runtime:                 RuntimeKimi,
		Supported:               true,
		Strategy:                StrategyKimiSessionTurns,
		RequiresPreparedPoint:   true,
		SupportsDetachedRestore: true,
		SupportsLiveRestore:     false,
	},
	RuntimePi: {
		Runtime:                 RuntimePi,
		Supported:               true,
		Strategy:                StrategyPiSessionTreeRewind,
		RequiresPreparedPoint:   true,
		SupportsDetachedRestore: true,
		SupportsLiveRestore:     false,
	},
	"opencode": {
		Runtime:                 "opencode",
		Supported:               true,
		Strategy:                StrategyOpenCodeSessionRevert,
		RequiresPreparedPoint:   true,
		SupportsDetachedRestore: true,
		SupportsLiveRestore:     true,
	},
}

// ForRuntime returns the recovery Capability for the given runtime name.
func ForRuntime(runtime string) Capability {
	trimmed := strings.TrimSpace(runtime)
	if trimmed == RuntimeBuiltin {
		capability := capabilityByRuntime[RuntimeOpenCode]
		capability.Runtime = trimmed
		return capability
	}
	if capability, ok := capabilityByRuntime[trimmed]; ok {
		return capability
	}
	return Capability{
		Runtime:           trimmed,
		UnsupportedReason: "unknown runtime",
	}
}

// SupportedRuntimes returns the sorted list of runtime names that support recovery.
func SupportedRuntimes() []string {
	supported := make([]string, 0, len(capabilityByRuntime))
	for runtime, capability := range capabilityByRuntime {
		if !capability.Supported {
			continue
		}
		supported = append(supported, runtime)
	}
	sort.Strings(supported)
	return supported
}

// EncodeCodexAppServerRecoveryRef serializes a codex app-server turn rollback
// recovery ref.
func EncodeCodexAppServerRecoveryRef(threadID string, turnCount int) (string, error) {
	refBytes, err := json.Marshal(CodexAppServerRecoveryRef{
		ThreadID:  strings.TrimSpace(threadID),
		TurnCount: turnCount,
	})
	if err != nil {
		return "", err
	}
	return string(refBytes), nil
}

// DecodeCodexAppServerRecoveryRef parses a codex app-server turn rollback
// recovery ref.
func DecodeCodexAppServerRecoveryRef(recoveryRef string) (CodexAppServerRecoveryRef, error) {
	var ref CodexAppServerRecoveryRef
	err := json.Unmarshal([]byte(recoveryRef), &ref)
	ref.ThreadID = strings.TrimSpace(ref.ThreadID)
	return ref, err
}

// EncodeClaudeSessionResumeAtRecoveryRef serializes a Claude Code resume-at
// recovery ref.
func EncodeClaudeSessionResumeAtRecoveryRef(sessionID string, boundaryUUID string, boundaryKind string) (string, error) {
	kind := strings.TrimSpace(boundaryKind)
	if kind == "" {
		kind = RecoveryBoundaryResumeAt
	}
	refBytes, err := json.Marshal(ClaudeSessionResumeAtRecoveryRef{
		SessionID:    strings.TrimSpace(sessionID),
		BoundaryUUID: strings.TrimSpace(boundaryUUID),
		BoundaryKind: kind,
	})
	if err != nil {
		return "", err
	}
	return string(refBytes), nil
}

// DecodeClaudeSessionResumeAtRecoveryRef parses a Claude Code resume-at
// recovery ref.
func DecodeClaudeSessionResumeAtRecoveryRef(recoveryRef string) (ClaudeSessionResumeAtRecoveryRef, error) {
	var ref ClaudeSessionResumeAtRecoveryRef
	err := json.Unmarshal([]byte(recoveryRef), &ref)
	ref.SessionID = strings.TrimSpace(ref.SessionID)
	ref.BoundaryUUID = strings.TrimSpace(ref.BoundaryUUID)
	ref.BoundaryKind = strings.TrimSpace(ref.BoundaryKind)
	if ref.BoundaryKind == "" {
		ref.BoundaryKind = RecoveryBoundaryResumeAt
	}
	return ref, err
}

// EncodeOpenCodeSessionRevertRecoveryRef serializes an OpenCode session revert
// recovery ref.
func EncodeOpenCodeSessionRevertRecoveryRef(sessionID string, messageID string) (string, error) {
	refBytes, err := json.Marshal(OpenCodeSessionRevertRecoveryRef{
		SessionID: strings.TrimSpace(sessionID),
		MessageID: strings.TrimSpace(messageID),
	})
	if err != nil {
		return "", err
	}
	return string(refBytes), nil
}

// DecodeOpenCodeSessionRevertRecoveryRef parses an OpenCode session revert
// recovery ref.
func DecodeOpenCodeSessionRevertRecoveryRef(recoveryRef string) (OpenCodeSessionRevertRecoveryRef, error) {
	var ref OpenCodeSessionRevertRecoveryRef
	err := json.Unmarshal([]byte(recoveryRef), &ref)
	ref.SessionID = strings.TrimSpace(ref.SessionID)
	ref.MessageID = strings.TrimSpace(ref.MessageID)
	return ref, err
}

// EncodeKimiSessionTurnsRecoveryRef serializes a Kimi completed-turn boundary.
func EncodeKimiSessionTurnsRecoveryRef(sessionID string, turnCount int, inputCounts ...int) (string, error) {
	inputCount := turnCount
	if len(inputCounts) > 0 {
		inputCount = inputCounts[0]
	}
	refBytes, err := json.Marshal(KimiSessionTurnsRecoveryRef{
		SessionID:  strings.TrimSpace(sessionID),
		TurnCount:  turnCount,
		InputCount: inputCount,
	})
	if err != nil {
		return "", err
	}
	return string(refBytes), nil
}

// DecodeKimiSessionTurnsRecoveryRef parses a Kimi completed-turn boundary.
func DecodeKimiSessionTurnsRecoveryRef(recoveryRef string) (KimiSessionTurnsRecoveryRef, error) {
	var ref KimiSessionTurnsRecoveryRef
	err := json.Unmarshal([]byte(recoveryRef), &ref)
	ref.SessionID = strings.TrimSpace(ref.SessionID)
	if ref.InputCount == 0 && ref.TurnCount > 0 && !strings.Contains(recoveryRef, `"inputCount"`) {
		ref.InputCount = ref.TurnCount
	}
	return ref, err
}

// EncodePiSessionTreeRewindRef serializes a Pi session-tree rewind ref.
func EncodePiSessionTreeRewindRef(sessionID string, boundaryEntryID string, boundaryKind string) (string, error) {
	kind := strings.TrimSpace(boundaryKind)
	if kind == "" {
		kind = RecoveryBoundaryResumeAt
	}
	refBytes, err := json.Marshal(PiSessionTreeRewindRef{
		SessionID:       strings.TrimSpace(sessionID),
		BoundaryEntryID: strings.TrimSpace(boundaryEntryID),
		BoundaryKind:    kind,
	})
	if err != nil {
		return "", err
	}
	return string(refBytes), nil
}

// DecodePiSessionTreeRewindRef parses a Pi session-tree rewind ref.
func DecodePiSessionTreeRewindRef(recoveryRef string) (PiSessionTreeRewindRef, error) {
	var ref PiSessionTreeRewindRef
	err := json.Unmarshal([]byte(recoveryRef), &ref)
	ref.SessionID = strings.TrimSpace(ref.SessionID)
	ref.BoundaryEntryID = strings.TrimSpace(ref.BoundaryEntryID)
	ref.BoundaryKind = strings.TrimSpace(ref.BoundaryKind)
	if ref.BoundaryKind == "" {
		ref.BoundaryKind = RecoveryBoundaryResumeAt
	}
	return ref, err
}

// ValidatePreparedRecoveryPoint validates that a prepared recovery point's
// strategy, threadID, and recoveryRef are consistent.
func ValidatePreparedRecoveryPoint(strategy string, threadID string, recoveryRef string) error {
	trimmedStrategy := strings.TrimSpace(strategy)
	trimmedThreadID := strings.TrimSpace(threadID)
	trimmedRecoveryRef := strings.TrimSpace(recoveryRef)
	if trimmedStrategy == "" {
		return fmt.Errorf("recovery strategy is required")
	}
	if trimmedThreadID == "" {
		return fmt.Errorf("recovery thread id is required")
	}
	if trimmedRecoveryRef == "" {
		return fmt.Errorf("recovery ref is required")
	}

	switch trimmedStrategy {
	case StrategyCodexAppServerTurns:
		ref, err := DecodeCodexAppServerRecoveryRef(trimmedRecoveryRef)
		if err != nil {
			return fmt.Errorf("decode codex recovery ref: %w", err)
		}
		if ref.ThreadID == "" {
			return fmt.Errorf("codex recovery ref is missing thread id")
		}
		if ref.TurnCount < 0 {
			return fmt.Errorf("codex recovery ref has invalid turn count")
		}
		if ref.ThreadID != trimmedThreadID {
			return fmt.Errorf("codex recovery ref thread id %q does not match thread id %q", ref.ThreadID, trimmedThreadID)
		}
		return nil
	case StrategyClaudeSessionResumeAt:
		ref, err := DecodeClaudeSessionResumeAtRecoveryRef(trimmedRecoveryRef)
		if err != nil {
			return fmt.Errorf("decode claude recovery ref: %w", err)
		}
		if ref.SessionID == "" {
			return fmt.Errorf("claude recovery ref is missing session id")
		}
		if ref.SessionID != trimmedThreadID {
			return fmt.Errorf("claude recovery ref session id %q does not match thread id %q", ref.SessionID, trimmedThreadID)
		}
		switch ref.BoundaryKind {
		case RecoveryBoundaryResumeAt:
			if ref.BoundaryUUID == "" {
				return fmt.Errorf("claude recovery ref is missing boundary uuid")
			}
		case RecoveryBoundaryBeforeFirstInput:
			// A zero-boundary restore creates a fresh runtime session and therefore
			// intentionally has no previous transcript uuid.
		default:
			return fmt.Errorf("claude recovery ref has unknown boundary kind %q", ref.BoundaryKind)
		}
		return nil
	case StrategyOpenCodeSessionRevert:
		ref, err := DecodeOpenCodeSessionRevertRecoveryRef(trimmedRecoveryRef)
		if err != nil {
			return fmt.Errorf("decode opencode recovery ref: %w", err)
		}
		if ref.SessionID == "" {
			return fmt.Errorf("opencode recovery ref is missing session id")
		}
		if ref.MessageID == "" {
			return fmt.Errorf("opencode recovery ref is missing message id")
		}
		if !strings.HasPrefix(ref.MessageID, "msg") {
			return fmt.Errorf("opencode recovery ref message id %q does not start with msg", ref.MessageID)
		}
		if ref.SessionID != trimmedThreadID {
			return fmt.Errorf("opencode recovery ref session id %q does not match thread id %q", ref.SessionID, trimmedThreadID)
		}
		return nil
	case StrategyKimiSessionTurns:
		ref, err := DecodeKimiSessionTurnsRecoveryRef(trimmedRecoveryRef)
		if err != nil {
			return fmt.Errorf("decode kimi recovery ref: %w", err)
		}
		if ref.SessionID == "" {
			return fmt.Errorf("kimi recovery ref is missing session id")
		}
		if ref.TurnCount < 0 {
			return fmt.Errorf("kimi recovery ref has invalid turn count")
		}
		if ref.InputCount < 0 {
			return fmt.Errorf("kimi recovery ref has invalid input count")
		}
		if ref.SessionID != trimmedThreadID {
			return fmt.Errorf("kimi recovery ref session id %q does not match thread id %q", ref.SessionID, trimmedThreadID)
		}
		return nil
	case StrategyPiSessionTreeRewind:
		ref, err := DecodePiSessionTreeRewindRef(trimmedRecoveryRef)
		if err != nil {
			return fmt.Errorf("decode pi recovery ref: %w", err)
		}
		if ref.SessionID == "" {
			return fmt.Errorf("pi recovery ref is missing session id")
		}
		if ref.SessionID != trimmedThreadID {
			return fmt.Errorf("pi recovery ref session id %q does not match thread id %q", ref.SessionID, trimmedThreadID)
		}
		switch ref.BoundaryKind {
		case RecoveryBoundaryResumeAt:
			if ref.BoundaryEntryID == "" {
				return fmt.Errorf("pi recovery ref is missing boundary entry id")
			}
		case RecoveryBoundaryBeforeFirstInput:
			// A zero-boundary restore creates a fresh runtime session and therefore
			// intentionally has no previous entry id.
		default:
			return fmt.Errorf("pi recovery ref has unknown boundary kind %q", ref.BoundaryKind)
		}
		return nil
	default:
		return fmt.Errorf("unknown recovery strategy %q", trimmedStrategy)
	}
}