package daemoncontrol

import (
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/x/proc"
)

var (
	ErrControlUnauthorized = errors.New("daemon control unauthorized")
	ErrControlConflict     = errors.New("daemon control conflict")
	ErrOwnerActive         = errors.New("daemon owner is still active")
	ErrOwnerUnsupported    = errors.New("daemon owner cannot be adopted")
	ErrOwnerIdentity       = errors.New("daemon owner identity is invalid")
)

// ReceiptRewriter is the durable compare-and-swap boundary used by ownership
// rotation. The old in-memory capability remains authoritative unless this
// rewrite succeeds.
type ReceiptRewriter interface {
	RewriteMatching(expected, replacement Receipt) error
}

// Credentials prove the caller owns the current receipt capability and owner
// process lifetime. Identity probes deliberately require only the token; every
// mutating operation requires this complete tuple.
type Credentials struct {
	ControlToken string
	InstanceID   string
	OwnerID      string
	Owner        ProcessIdentity
}

type Identity struct {
	InstanceID        string    `json:"instanceId"`
	DesktopAPIVersion int       `json:"desktopApiVersion"`
	PID               int       `json:"pid"`
	StartedAt         time.Time `json:"processStartedAt"`
	BaseURL           string    `json:"baseURL"`
	OwnerKind         OwnerKind `json:"ownerKind"`
	Draining          bool      `json:"draining"`
}

// Adoption rotates a dead Electron owner's identity and the control
// capability as one durable operation.
type Adoption struct {
	Owner        Owner
	ControlToken string
}

type ControllerOptions struct {
	Receipt       Receipt
	Store         ReceiptRewriter
	LookupProcess func(int) (proc.ProcessInfo, bool)
	PIDAlive      func(int) bool
	NewToken      func() (string, error)
}

// Controller serializes every control-plane mutation and owns the current
// receipt capability.
type Controller struct {
	mu            sync.Mutex
	receipt       Receipt
	store         ReceiptRewriter
	lookupProcess func(int) (proc.ProcessInfo, bool)
	pidAlive      func(int) bool
	newToken      func() (string, error)
	draining      bool
}

func NewController(options ControllerOptions) (*Controller, error) {
	receipt := options.Receipt
	if err := receipt.Validate(); err != nil {
		return nil, err
	}
	if options.Store == nil {
		return nil, fmt.Errorf("daemon control receipt store is required")
	}
	lookup := options.LookupProcess
	if lookup == nil {
		lookup = proc.ProcessByPID
	}
	pidAlive := options.PIDAlive
	if pidAlive == nil {
		pidAlive = proc.PIDAlive
	}
	newToken := options.NewToken
	if newToken == nil {
		newToken = newControlToken
	}
	return &Controller{
		receipt:       receipt,
		store:         options.Store,
		lookupProcess: lookup,
		pidAlive:      pidAlive,
		newToken:      newToken,
	}, nil
}

func (c *Controller) Identity(controlToken string) (Identity, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if !sameSecret(controlToken, c.receipt.ControlToken) {
		return Identity{}, ErrControlUnauthorized
	}
	return c.identityLocked(), nil
}

func (c *Controller) PublicIdentity() Identity {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.identityLocked()
}

func (c *Controller) CurrentReceipt() Receipt {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.receipt
}

func (c *Controller) Authorize(credentials Credentials) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.authorizeLocked(credentials)
}

func (c *Controller) Adopt(credentials Credentials, next Owner) (Adoption, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if err := c.authorizeStoredLocked(credentials); err != nil {
		return Adoption{}, err
	}
	if c.draining {
		return Adoption{}, fmt.Errorf("%w: daemon is draining", ErrControlConflict)
	}
	if c.receipt.Owner.Kind != OwnerKindElectron {
		return Adoption{}, ErrOwnerUnsupported
	}
	active, ambiguous := c.processLifetimeState(c.receipt.Owner.ProcessIdentity)
	if active {
		return Adoption{}, ErrOwnerActive
	}
	if ambiguous {
		return Adoption{}, fmt.Errorf("%w: previous owner process lifetime is ambiguous", ErrOwnerIdentity)
	}
	if next.Kind != OwnerKindElectron || strings.TrimSpace(next.ID) == "" {
		return Adoption{}, ErrOwnerIdentity
	}
	actual, found := c.lookupProcess(next.PID)
	if !found || !next.ProcessIdentity.Matches(actual) {
		return Adoption{}, ErrOwnerIdentity
	}
	token, err := c.newToken()
	if err != nil {
		return Adoption{}, fmt.Errorf("generate rotated control token: %w", err)
	}
	next.ID = strings.TrimSpace(next.ID)
	replacement := c.receipt
	replacement.Owner = next
	replacement.ControlToken = token
	if err := c.store.RewriteMatching(c.receipt, replacement); err != nil {
		return Adoption{}, fmt.Errorf("rewrite adopted daemon receipt: %w", err)
	}
	c.receipt = replacement
	return Adoption{Owner: next, ControlToken: token}, nil
}

// BeginShutdown verifies the current owner and transitions the controller to
// draining. The caller schedules the daemon drain only after committing its
// HTTP response.
func (c *Controller) BeginShutdown(credentials Credentials) (alreadyDraining bool, err error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if err := c.authorizeLocked(credentials); err != nil {
		return false, err
	}
	if c.receipt.Owner.Kind != OwnerKindElectron {
		return false, ErrOwnerUnsupported
	}
	if c.draining {
		return true, nil
	}
	c.draining = true
	return false, nil
}

func (c *Controller) authorizeLocked(credentials Credentials) error {
	if err := c.authorizeStoredLocked(credentials); err != nil {
		return err
	}
	if !c.processLifetimeMatches(c.receipt.Owner.ProcessIdentity) {
		return ErrOwnerIdentity
	}
	return nil
}

func (c *Controller) authorizeStoredLocked(credentials Credentials) error {
	if !sameSecret(credentials.ControlToken, c.receipt.ControlToken) ||
		strings.TrimSpace(credentials.InstanceID) != c.receipt.InstanceID ||
		strings.TrimSpace(credentials.OwnerID) != c.receipt.Owner.ID ||
		!sameIdentity(credentials.Owner, c.receipt.Owner.ProcessIdentity) {
		return ErrControlUnauthorized
	}
	return nil
}

func (c *Controller) processLifetimeMatches(identity ProcessIdentity) bool {
	actual, found := c.lookupProcess(identity.PID)
	if !found {
		return false
	}
	return identity.Matches(actual)
}

func (c *Controller) processLifetimeState(identity ProcessIdentity) (active, ambiguous bool) {
	actual, found := c.lookupProcess(identity.PID)
	if found {
		return identity.Matches(actual), false
	}
	if c.pidAlive(identity.PID) {
		return false, true
	}
	return false, false
}

func (c *Controller) identityLocked() Identity {
	return Identity{
		InstanceID:        c.receipt.InstanceID,
		DesktopAPIVersion: c.receipt.DesktopAPIVersion,
		PID:               c.receipt.PID,
		StartedAt:         c.receipt.StartedAt,
		BaseURL:           c.receipt.BaseURL,
		OwnerKind:         c.receipt.Owner.Kind,
		Draining:          c.draining,
	}
}

func sameSecret(provided, expected string) bool {
	provided = strings.TrimSpace(provided)
	expected = strings.TrimSpace(expected)
	if provided == "" || expected == "" || len(provided) != len(expected) {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(provided), []byte(expected)) == 1
}

func newControlToken() (string, error) {
	data := make([]byte, 32)
	if _, err := rand.Read(data); err != nil {
		return "", err
	}
	return base64.RawURLEncoding.EncodeToString(data), nil
}