package store

import (
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"database/sql"
	"errors"
	"fmt"
	"time"

	"github.com/klauspost/compress/zstd"
	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
	"github.com/rvbox/rvbox/internal/domain"
	"google.golang.org/protobuf/proto"
)

var (
	ErrCommandConflict = errors.New("command UUID conflicts with immutable request")
	ErrClientNotFound  = errors.New("target client does not exist")
)

// QueueCommandInput is the already-validated immutable command payload. The
// control boundary validates and serializes the ExecutionSpec before it enters
// this storage layer; the store owns only durable admission and quota charging.
type QueueCommandInput struct {
	IssueUUID       domain.UUID
	ClientID        string
	IssueTime       time.Time
	ReceiptTime     time.Time
	QueueExpiryTime *time.Time
	ImmutableSHA256 [32]byte
	ExecutionSpec   []byte
	// ScriptPresent distinguishes a valid zero-byte script from a command-text
	// request with no script payload. When true, ScriptContent is checked
	// against the descriptor embedded in ExecutionSpec and stored as a
	// command-owned compressed payload.
	ScriptPresent bool
	ScriptContent []byte
}

type QueueCommandResult struct {
	Duplicate bool
	Lifecycle uint32
	Revision  uint64
}

// QueueCommand atomically persists a queued command before it is visible to a
// dispatcher. Repeating the same UUID/hash is idempotent; a different hash is
// a hard conflict and cannot replace a previously admitted command.
func (store *Store) QueueCommand(ctx context.Context, input QueueCommandInput) (QueueCommandResult, error) {
	if !validCommandInput(input) {
		return QueueCommandResult{}, errors.New("invalid queued command")
	}
	stored, err := compressCommandSpec(input.ExecutionSpec)
	if err != nil {
		return QueueCommandResult{}, err
	}
	scriptStored, scriptDescriptor, err := validateAndCompressScript(input)
	if err != nil {
		return QueueCommandResult{}, err
	}
	rows, indexes := uint64(1), uint64(2)
	encodedBytes := uint64(len(stored))
	if input.ScriptPresent {
		rows++
		indexes++
		encodedBytes += uint64(len(scriptStored))
	}
	charge, err := EstimateCharge(ChargeInput{EncodedBytes: encodedBytes, SQLiteRows: rows, IndexEntries: indexes})
	if err != nil {
		return QueueCommandResult{}, err
	}

	store.writeMu.Lock()
	defer store.writeMu.Unlock()
	database, err := store.openDatabase()
	if err != nil {
		return QueueCommandResult{}, err
	}
	tx, err := database.BeginTx(ctx, nil)
	if err != nil {
		return QueueCommandResult{}, err
	}
	defer tx.Rollback()

	var storedHash []byte
	var lifecycle uint32
	var revision uint64
	err = tx.QueryRowContext(ctx, `SELECT immutable_request_sha256, lifecycle, revision FROM commands WHERE issue_uuid = ?`, input.IssueUUID[:]).Scan(&storedHash, &lifecycle, &revision)
	if err == nil {
		if len(storedHash) != len(input.ImmutableSHA256) || subtle.ConstantTimeCompare(storedHash, input.ImmutableSHA256[:]) != 1 {
			return QueueCommandResult{}, ErrCommandConflict
		}
		return QueueCommandResult{Duplicate: true, Lifecycle: lifecycle, Revision: revision}, tx.Commit()
	}
	if !errors.Is(err, sql.ErrNoRows) {
		return QueueCommandResult{}, err
	}

	var clientCharged, serverCharged uint64
	err = tx.QueryRowContext(ctx, `SELECT clients.charged_bytes, storage_counters.command_charged_bytes
FROM clients JOIN storage_counters ON storage_counters.singleton = 1 WHERE clients.client_id = ?`, input.ClientID).Scan(&clientCharged, &serverCharged)
	if errors.Is(err, sql.ErrNoRows) {
		return QueueCommandResult{}, ErrClientNotFound
	}
	if err != nil {
		return QueueCommandResult{}, err
	}
	freeBytes, err := store.freeSpaceProbe.AvailableBytes(store.dataDir)
	if err != nil {
		return QueueCommandResult{}, fmt.Errorf("check filesystem capacity: %w", err)
	}
	decision, err := CheckReservation(store.quotaLimits, ReservationState{
		ClientTotalCharged: clientCharged, ServerTotalCharged: serverCharged,
		CloseoutRemaining: store.quotaLimits.CloseoutReserveBytes, FilesystemFreeBytes: freeBytes,
	}, ReservationRequest{ChargedBytes: charge, PhysicalBytes: encodedBytes})
	if err != nil {
		return QueueCommandResult{}, err
	}
	var expiry any
	if input.QueueExpiryTime != nil {
		expiry = input.QueueExpiryTime.UTC().UnixNano()
	}
	_, err = tx.ExecContext(ctx, `INSERT INTO commands (
issue_uuid, client_id, issue_time, server_receipt_time, queue_expiry_time,
lifecycle, revision, retained_compressed_bytes, charged_bytes,
closeout_remaining_bytes, immutable_request_sha256, execution_spec,
execution_spec_raw_bytes, execution_spec_stored_bytes, execution_spec_compression
) VALUES (?, ?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, 2)`,
		input.IssueUUID[:], input.ClientID, input.IssueTime.UTC().UnixNano(), input.ReceiptTime.UTC().UnixNano(), expiry,
		encodedBytes, charge, decision.CloseoutRemaining, input.ImmutableSHA256[:], stored, len(input.ExecutionSpec), len(stored))
	if err != nil {
		return QueueCommandResult{}, err
	}
	if input.ScriptPresent {
		if _, err := tx.ExecContext(ctx, `INSERT INTO command_payloads (
issue_uuid, kind, raw_bytes, stored_bytes, compression, sha256, inline_data, segment_path
) VALUES (?, 'script', ?, ?, 2, ?, ?, NULL)`, input.IssueUUID[:], scriptDescriptor.GetSizeBytes(), len(scriptStored), scriptDescriptor.GetSha256(), scriptStored); err != nil {
			return QueueCommandResult{}, err
		}
	}
	if _, err := tx.ExecContext(ctx, `UPDATE clients SET charged_bytes = ? WHERE client_id = ?`, decision.ClientTotalCharged, input.ClientID); err != nil {
		return QueueCommandResult{}, err
	}
	if _, err := tx.ExecContext(ctx, `UPDATE storage_counters SET command_charged_bytes = ? WHERE singleton = 1`, decision.ServerTotalCharged); err != nil {
		return QueueCommandResult{}, err
	}
	if err := tx.Commit(); err != nil {
		return QueueCommandResult{}, err
	}
	return QueueCommandResult{Lifecycle: 1, Revision: 1}, nil
}

func validCommandInput(input QueueCommandInput) bool {
	if isZeroUUID([16]byte(input.IssueUUID)) || input.ClientID == "" || len(input.ClientID) > 128 || input.IssueTime.IsZero() || input.ReceiptTime.IsZero() || len(input.ExecutionSpec) == 0 || allZero(input.ImmutableSHA256[:]) {
		return false
	}
	if !input.ScriptPresent && len(input.ScriptContent) != 0 {
		return false
	}
	return input.QueueExpiryTime == nil || input.QueueExpiryTime.After(input.IssueTime)
}

const maxStoredScriptBytes = 10 << 20

func validateAndCompressScript(input QueueCommandInput) ([]byte, *rvboxv1.ScriptDescriptor, error) {
	if !input.ScriptPresent {
		return nil, nil, nil
	}
	var spec rvboxv1.ExecutionSpec
	if err := proto.Unmarshal(input.ExecutionSpec, &spec); err != nil {
		return nil, nil, fmt.Errorf("decode execution spec for script payload: %w", err)
	}
	descriptor := spec.GetScript()
	if descriptor == nil || descriptor.GetSizeBytes() > maxStoredScriptBytes || len(descriptor.GetSha256()) != sha256.Size || uint64(len(input.ScriptContent)) != descriptor.GetSizeBytes() {
		return nil, nil, errors.New("script payload does not match execution descriptor")
	}
	digest := sha256.Sum256(input.ScriptContent)
	if !bytes.Equal(digest[:], descriptor.GetSha256()) {
		return nil, nil, errors.New("script payload digest does not match execution descriptor")
	}
	stored, err := compressCommandSpec(input.ScriptContent)
	if err != nil {
		return nil, nil, err
	}
	return stored, descriptor, nil
}

func compressCommandSpec(spec []byte) ([]byte, error) {
	encoder, err := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1))
	if err != nil {
		return nil, err
	}
	defer encoder.Close()
	return encoder.EncodeAll(spec, nil), nil
}