package store

import (
	"bytes"
	"context"
	"crypto/sha256"
	"database/sql"

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

// checkCommandPayloads verifies the inline compressed payload table and its
// relationship to script descriptors. A valid SQLite page is not enough: a
// short/changed zstd frame must make the owning scope dirty before dispatch.
func checkCommandPayloads(ctx context.Context, database *sql.DB) error {
	rows, err := database.QueryContext(ctx, `SELECT issue_uuid, kind, raw_bytes, stored_bytes, compression, sha256, inline_data, segment_path FROM command_payloads`)
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var issue, digest, stored []byte
		var kind string
		var rawBytes, storedBytes uint64
		var compression uint32
		var segmentPath sql.NullString
		if err := rows.Scan(&issue, &kind, &rawBytes, &storedBytes, &compression, &digest, &stored, &segmentPath); err != nil {
			return err
		}
		if len(issue) != 16 || kind != "script" || compression != 2 || rawBytes > maxStoredScriptBytes || storedBytes != uint64(len(stored)) || len(digest) != sha256.Size || segmentPath.Valid || len(stored) == 0 {
			return ErrCommittedRangeMissing
		}
		decoded, err := decodeCommandPayload(stored, rawBytes, maxStoredScriptBytes)
		if err != nil {
			return err
		}
		computed := sha256.Sum256(decoded)
		if !bytes.Equal(computed[:], digest) {
			return ErrCommittedRangeMissing
		}
	}
	if err := rows.Err(); err != nil {
		return err
	}

	commands, err := database.QueryContext(ctx, `SELECT issue_uuid, execution_spec FROM commands`)
	if err != nil {
		return err
	}
	defer commands.Close()
	for commands.Next() {
		var issue, encoded []byte
		if err := commands.Scan(&issue, &encoded); err != nil {
			return err
		}
		if len(issue) != 16 {
			return ErrCommittedRangeMissing
		}
		var spec rvboxv1.ExecutionSpec
		if err := proto.Unmarshal(encoded, &spec); err != nil {
			// Some low-level store tests intentionally use opaque bytes. Such
			// records cannot claim a script source, and dispatch validation will
			// reject them before any user code runs.
			continue
		}
		descriptor := spec.GetScript()
		var payloadCount uint64
		if err := database.QueryRowContext(ctx, `SELECT count(*) FROM command_payloads WHERE issue_uuid = ? AND kind = 'script'`, issue).Scan(&payloadCount); err != nil {
			return err
		}
		if descriptor == nil {
			if payloadCount != 0 {
				return ErrCommittedRangeMissing
			}
			continue
		}
		if payloadCount != 1 || descriptor.GetSizeBytes() > maxStoredScriptBytes || len(descriptor.GetSha256()) != sha256.Size {
			return ErrCommittedRangeMissing
		}
		var rawBytes, storedBytes uint64
		var compression uint32
		var digest, stored []byte
		if err := database.QueryRowContext(ctx, `SELECT raw_bytes, stored_bytes, compression, sha256, inline_data FROM command_payloads WHERE issue_uuid = ? AND kind = 'script'`, issue).Scan(&rawBytes, &storedBytes, &compression, &digest, &stored); err != nil {
			return err
		}
		if rawBytes != descriptor.GetSizeBytes() || compression != 2 || len(digest) != sha256.Size || !bytes.Equal(digest, descriptor.GetSha256()) {
			return ErrCommittedRangeMissing
		}
		decoded, err := decodeCommandPayload(stored, rawBytes, maxStoredScriptBytes)
		computed := sha256.Sum256(decoded)
		if err != nil || !bytes.Equal(computed[:], descriptor.GetSha256()) {
			return ErrCommittedRangeMissing
		}
	}
	return commands.Err()
}

func decodeCommandPayload(stored []byte, expected, maximum uint64) ([]byte, error) {
	if expected > maximum {
		return nil, ErrCommittedRangeMissing
	}
	decoder, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(maximum+1))
	if err != nil {
		return nil, ErrCommittedRangeMissing
	}
	defer decoder.Close()
	decoded, err := decoder.DecodeAll(stored, nil)
	if err != nil || uint64(len(decoded)) != expected || uint64(len(decoded)) > maximum {
		return nil, ErrCommittedRangeMissing
	}
	return decoded, nil
}