package mysql

import (
	"context"
	"crypto/subtle"
	"database/sql"
	"encoding/json"
	"errors"
	"time"

	"git.yuntuo.hk/k2safe/orbit/internal/fault"
	core "git.yuntuo.hk/k2safe/orbit/model"
	"git.yuntuo.hk/k2safe/orbit/platform/auth"
	"git.yuntuo.hk/k2safe/orbit/pulse/worker"
)

// ExecutionOptions 是服务端执行窗口，不由 Worker 任意指定。
type ExecutionOptions struct {
	LeaseDuration time.Duration // LeaseDuration 每次领取/续租的最大窗口。
	SessionTTL    time.Duration // SessionTTL 会话心跳有效期，不代替执行租约。
}

// ExecutionStore 所有执行权以 MySQL 为准，节点内存不持有可绕过数据库的授权。
type ExecutionStore struct {
	db      *sql.DB
	node    core.ID
	options ExecutionOptions
}

var _ worker.Store = (*ExecutionStore)(nil)

func (s *ExecutionStore) Authorize(ctx context.Context, token auth.Token) error {
	return s.transaction(ctx, token, "workers:connect", func(*sql.Tx, auth.Permissions) error { return nil })
}

func NewExecutionStore(db *sql.DB, node core.ID, options ExecutionOptions) (*ExecutionStore, error) {
	if options.LeaseDuration == 0 {
		options.LeaseDuration = 30 * time.Second
	}
	if options.SessionTTL == 0 {
		options.SessionTTL = 45 * time.Second
	}
	if node.Validate() != nil || options.LeaseDuration < time.Second || options.LeaseDuration > 2*time.Minute || options.SessionTTL < options.LeaseDuration || options.SessionTTL > 5*time.Minute {
		return nil, fault.Invalid("invalid execution windows or node identity")
	}
	return &ExecutionStore{db: db, node: node, options: options}, nil
}
func (s *ExecutionStore) transaction(ctx context.Context, token auth.Token, action string, fn func(*sql.Tx, auth.Permissions) error) error {
	return (&TaskStore{db: s.db}).transaction(ctx, token, action, fn)
}

type sessionState struct {
	queue        core.ID
	capacity     uint32
	sequence     uint64
	draining     bool
	expires      time.Time
	capabilities []worker.Capability
}

func (s *ExecutionStore) welcome(b worker.Binding, state sessionState, now time.Time) worker.Welcome {
	return worker.Welcome{SessionID: b.SessionID, Epoch: b.Epoch, ServerTime: now, ExpiresAt: state.expires, LeaseDurationMS: uint64(s.options.LeaseDuration / time.Millisecond), Draining: state.draining}
}
func (s *ExecutionStore) Connect(ctx context.Context, token auth.Token, in worker.ConnectInput) (worker.Binding, worker.Welcome, error) {
	var binding worker.Binding
	var out worker.Welcome
	proof, err := worker.SecretDigest(in.InstanceSecret)
	if err != nil {
		return binding, out, err
	}
	registration, err := requestDigest(struct {
		Queue        core.ID
		Capabilities []worker.Capability
	}{in.QueueID, in.Capabilities})
	if err != nil {
		return binding, out, err
	}
	err = s.transaction(ctx, token, "workers:connect", func(tx *sql.Tx, p auth.Permissions) error {
		if !p.AllowsQueue(in.QueueID) {
			return fault.Forbidden
		}
		_, err := tx.ExecContext(ctx, `INSERT INTO pulse_worker_instances(namespace_id,id,credential_id,queue_id,secret_digest,registration_digest) VALUES (?,?,?,?,?,?) ON DUPLICATE KEY UPDATE id=pulse_worker_instances.id`, token.NamespaceID, in.InstanceID, token.CredentialID, in.QueueID, proof, registration)
		if err != nil {
			return err
		}
		var credential core.ID
		var storedProof, storedRegistration core.Digest
		var epoch uint64
		var old *core.ID
		var draining bool
		err = tx.QueryRowContext(ctx, `SELECT credential_id,secret_digest,registration_digest,session_epoch,current_session_id,draining FROM pulse_worker_instances WHERE namespace_id=? AND id=? FOR UPDATE`, token.NamespaceID, in.InstanceID).Scan(&credential, &storedProof, &storedRegistration, &epoch, &old, &draining)
		if err != nil {
			return err
		}
		if credential != token.CredentialID || subtle.ConstantTimeCompare(proof[:], storedProof[:]) != 1 {
			return fault.Forbidden
		}
		if registration != storedRegistration || epoch == ^uint64(0) {
			return fault.Conflict
		}
		now, err := dbNow(ctx, tx)
		if err != nil {
			return err
		}
		session, err := core.NewID()
		if err != nil {
			return err
		}
		binding = worker.Binding{InstanceID: in.InstanceID, SessionID: session, Epoch: epoch + 1, Proof: proof}
		state := sessionState{queue: in.QueueID, capacity: in.Capacity, draining: draining, expires: now.Add(s.options.SessionTTL), capabilities: in.Capabilities}
		caps, err := json.Marshal(in.Capabilities)
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, `INSERT INTO pulse_worker_sessions(namespace_id,id,worker_instance_id,credential_id,node_id,session_epoch,capabilities,capacity,draining,expires_at) VALUES (?,?,?,?,?,?,?,?,?,?)`, token.NamespaceID, session, in.InstanceID, token.CredentialID, s.node, binding.Epoch, string(caps), in.Capacity, draining, state.expires)
		if err != nil {
			return err
		}
		if old != nil {
			if _, err := tx.ExecContext(ctx, `UPDATE pulse_worker_sessions SET expires_at=?,draining=TRUE,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=?`, now, now, token.NamespaceID, *old); err != nil {
				return err
			}
		}
		_, err = tx.ExecContext(ctx, `UPDATE pulse_worker_instances SET current_session_id=?,session_epoch=?,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=?`, session, binding.Epoch, now, token.NamespaceID, in.InstanceID)
		if err != nil {
			return err
		}
		if err := emitEvent(ctx, tx, token.NamespaceID, session, 1, "worker_session", "worker.connected", map[string]any{"instance_id": in.InstanceID, "node_id": s.node, "epoch": binding.Epoch}); err != nil {
			return err
		}
		out = s.welcome(binding, state, now)
		return nil
	})
	if err != nil {
		return worker.Binding{}, worker.Welcome{}, err
	}
	return binding, out, nil
}

// lockSession 的顺序为凭据 -> 进程头 -> 当前会话；新连接和旧连接命令在同一行上串行化。
func (s *ExecutionStore) lockSession(ctx context.Context, tx *sql.Tx, token auth.Token, b worker.Binding, p auth.Permissions) (sessionState, error) {
	var state sessionState
	var credential core.ID
	var proof core.Digest
	var current *core.ID
	var epoch uint64
	if b.InstanceID.Validate() != nil || b.SessionID.Validate() != nil || b.Epoch == 0 {
		return state, worker.StaleSession
	}
	err := tx.QueryRowContext(ctx, `SELECT credential_id,queue_id,secret_digest,current_session_id,session_epoch,draining FROM pulse_worker_instances WHERE namespace_id=? AND id=? FOR UPDATE`, token.NamespaceID, b.InstanceID).Scan(&credential, &state.queue, &proof, &current, &epoch, &state.draining)
	if errors.Is(err, sql.ErrNoRows) {
		return state, worker.StaleSession
	}
	if err != nil {
		return state, err
	}
	if credential != token.CredentialID || current == nil || *current != b.SessionID || epoch != b.Epoch || subtle.ConstantTimeCompare(proof[:], b.Proof[:]) != 1 {
		return state, worker.StaleSession
	}
	if !p.AllowsQueue(state.queue) {
		return state, fault.Forbidden
	}
	var node core.ID
	var raw []byte
	err = tx.QueryRowContext(ctx, `SELECT node_id,capacity,capacity_sequence,capabilities,expires_at FROM pulse_worker_sessions WHERE namespace_id=? AND id=? FOR UPDATE`, token.NamespaceID, b.SessionID).Scan(&node, &state.capacity, &state.sequence, &raw, &state.expires)
	if err != nil {
		return state, err
	}
	now, err := dbNow(ctx, tx)
	if err != nil {
		return state, err
	}
	if node != s.node || !state.expires.After(now) {
		return state, worker.StaleSession
	}
	if err := json.Unmarshal(raw, &state.capabilities); err != nil {
		return state, err
	}
	return state, nil
}
func (s *ExecutionStore) Capacity(ctx context.Context, token auth.Token, b worker.Binding, in worker.CapacityInput) (worker.Welcome, error) {
	var out worker.Welcome
	err := s.transaction(ctx, token, "workers:consume", func(tx *sql.Tx, p auth.Permissions) error {
		state, err := s.lockSession(ctx, tx, token, b, p)
		if err != nil {
			return err
		}
		if in.Sequence < state.sequence || (in.Sequence == state.sequence && (in.Capacity != state.capacity || in.Drain != state.draining)) || (state.draining && !in.Drain) {
			return fault.Conflict
		}
		now, err := dbNow(ctx, tx)
		if err != nil {
			return err
		}
		state.expires = now.Add(s.options.SessionTTL)
		state.capacity = in.Capacity
		state.draining = state.draining || in.Drain
		_, err = tx.ExecContext(ctx, `UPDATE pulse_worker_sessions SET capacity=?,capacity_sequence=?,draining=?,expires_at=?,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=?`, state.capacity, in.Sequence, state.draining, state.expires, now, token.NamespaceID, b.SessionID)
		if err != nil {
			return err
		}
		if in.Drain {
			if _, err := tx.ExecContext(ctx, `UPDATE pulse_worker_instances SET draining=TRUE,revision=revision+1,updated_at=? WHERE namespace_id=? AND id=? AND draining=FALSE`, now, token.NamespaceID, b.InstanceID); err != nil {
				return err
			}
		}
		out = s.welcome(b, state, now)
		return nil
	})
	return out, err
}