package check

import (
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"sort"
	"strings"
	"time"
)

// The composition requests a check sends (issue 1921, notes 66-67): the request the checker sends the composer for one
// node, the checking program's whole input of a composition, is kept in the checker's own state RequestsStateName, so a
// session reproduces a check's composition exactly. The checker is a foundation part (novox/hq ADR 0619 §1-3), so its
// kept data is a section of its manifest's old form, here `state`, and not an object store asked in `requires`.
//
// **What bounds it.** A request is the node-engine's whole composition input of one node, 0.7 to 0.8 MB as JSON and
// growing with the mesh. A state's value holds at most 256 KiB (kvlimit.ModuleStateValueBytes, headers included), and
// the compressed request already passed it (issue 1966: 300127 bytes), so it is kept gzip-compressed in pieces: the key
// of the check and node holds a small document naming the request's digest, size and number of pieces, and each piece is
// a value of its own under RequestChunkKey, of at most RequestChunkBytes. The kept data is the checker's, a foundation
// part's, so it stays in the old section `state` (ADR 0619 §2) and is read through the tool runner's state read, which
// serves no object store of the old sections. The manifest's ttl-seconds is the lifetime (3 days) of the pointer and of
// every piece: the bus drops a value past it, so nothing waits on a timer here. A check kept again overwrites the key;
// a request whose pieces cannot all be put is refused naming the size (never dropped, never half kept: the pointer is
// put after its pieces).
const (
	// RequestsStateName is the checker's state of kept requests and RequestsBucket its bucket.
	RequestsStateName = "compositions"
	RequestsBucket    = CheckerModule + "_" + RequestsStateName
	// RequestsReadTool is the tool that reads a kept request.
	RequestsReadTool = "checker_composition_requests"
	// RequestChunkBytes is the most compressed bytes one piece holds: base64 of it, in its document, stays well under
	// the state's 256 KiB value.
	RequestChunkBytes = 128 << 10
)

// KeptRequest is the pointer to a request as the state holds it: whose it is, and what its pieces make.
type KeptRequest struct {
	V      int       `json:"v"`
	Check  string    `json:"check"`
	Node   string    `json:"node"`
	KeptAt time.Time `json:"kept_at"`
	// Size and Digest are the request's, uncompressed: its bytes and their sha256.
	Size   int    `json:"size"`
	Digest string `json:"digest"`
	// Pieces is how many pieces (RequestChunkKey of the digest, 0 to Pieces-1) make the request's gzip bytes.
	Pieces int `json:"pieces"`
}

// KeptPiece is one piece of a request's gzip bytes, base64-encoded.
type KeptPiece struct {
	V      int    `json:"v"`
	Digest string `json:"digest"`
	Index  int    `json:"index"`
	Gzip   string `json:"gzip"`
}

// isPieceKey is whether a key is a piece's: under pieces/ and a sha256 digest, which no check's key (group.heads/node) is.
func isPieceKey(k string) bool {
	rest, ok := strings.CutPrefix(k, "pieces/")
	if !ok {
		return false
	}
	digest, _, _ := strings.Cut(rest, "/")
	return len(digest) == 64
}

// RequestChunkKey is the key of a request's piece: under the request's digest, so a request kept again for the same
// bytes writes the same pieces.
func RequestChunkKey(digest string, index int) string {
	return fmt.Sprintf("pieces/%s/%d", digest, index)
}

// RequestKey is the key a check's request to a node is kept under: the check's key (group.heads) and the node.
func RequestKey(key, node string) string { return key + "/" + node }

// RequestPutter writes a value of the state, replacing one of that key.
type RequestPutter interface {
	Put(ctx context.Context, key string, value []byte) (uint64, error)
}

// RequestKeeper keeps the requests of a check.
type RequestKeeper struct {
	Store RequestPutter
	Now   func() time.Time
}

// Keep stores body under the check's key and the node. An error says why it could not, for the verdict's words.
func (k RequestKeeper) Keep(ctx context.Context, key, node string, body []byte) error {
	var packed bytes.Buffer
	zw := gzip.NewWriter(&packed)
	if _, err := zw.Write(body); err != nil {
		return fmt.Errorf("the composition request of check %s to %s could not be compressed: %w", key, node, err)
	}
	if err := zw.Close(); err != nil {
		return fmt.Errorf("the composition request of check %s to %s could not be compressed: %w", key, node, err)
	}
	now := time.Now
	if k.Now != nil {
		now = k.Now
	}
	sum := sha256.Sum256(body)
	digest := hex.EncodeToString(sum[:])
	pieces := (packed.Len() + RequestChunkBytes - 1) / RequestChunkBytes
	fail := func(err error) error {
		return fmt.Errorf("the composition request of check %s to %s (%d bytes, %d compressed in %d pieces) could not be kept in the state %s: %w",
			key, node, len(body), packed.Len(), pieces, RequestsStateName, err)
	}
	for i := 0; i < pieces; i++ {
		end := min((i+1)*RequestChunkBytes, packed.Len())
		doc, err := json.Marshal(KeptPiece{V: 2, Digest: digest, Index: i, Gzip: base64.StdEncoding.EncodeToString(packed.Bytes()[i*RequestChunkBytes : end])})
		if err != nil {
			return err
		}
		if _, err := k.Store.Put(ctx, RequestChunkKey(digest, i), doc); err != nil {
			return fail(err)
		}
	}
	doc, err := json.Marshal(KeptRequest{V: 2, Check: key, Node: node, KeptAt: now().UTC(), Size: len(body), Digest: digest, Pieces: pieces})
	if err != nil {
		return err
	}
	if _, err := k.Store.Put(ctx, RequestKey(key, node), doc); err != nil {
		return fail(err)
	}
	return nil
}

// ReadRequests answers the kept requests: with a check and a node, that request whole; with a check, the nodes it kept
// one for; with neither, every kept request. A node without a check is refused. get is the state's value of a key.
func ReadRequests(s StateRead, key, node string) (any, error) {
	if s.Keys == nil || s.Get == nil {
		return nil, fmt.Errorf("the state %s cannot be read here", RequestsStateName)
	}
	if node != "" && key == "" {
		return nil, fmt.Errorf("a node is asked for with the check's key (group.heads) of the request")
	}
	if node != "" {
		raw, err := s.Get(RequestKey(key, node))
		if err != nil {
			return nil, err
		}
		if len(raw) == 0 {
			return nil, fmt.Errorf("no composition request of check %s to %s is kept (none was sent, or it is past its lifetime)", key, node)
		}
		kept, body, err := unpackRequest(s, raw)
		if err != nil {
			return nil, fmt.Errorf("the kept request %s: %w", RequestKey(key, node), err)
		}
		return map[string]any{"check": kept.Check, "node": kept.Node, "size": kept.Size, "digest": kept.Digest, "kept_at": kept.KeptAt, "request": json.RawMessage(body)}, nil
	}
	keys, err := s.Keys()
	if err != nil {
		return nil, err
	}
	sort.Strings(keys)
	listed := []map[string]any{}
	for _, k := range keys {
		if key != "" && !strings.HasPrefix(k, key+"/") {
			continue
		}
		if isPieceKey(k) {
			continue // a piece of a request, not a request: never read here
		}
		raw, err := s.Get(k)
		if err != nil {
			return nil, err
		}
		var kept KeptRequest
		if err := json.Unmarshal(raw, &kept); err != nil {
			listed = append(listed, map[string]any{"key": k, "unreadable": err.Error()})
			continue
		}
		listed = append(listed, map[string]any{"check": kept.Check, "node": kept.Node, "size": kept.Size, "digest": kept.Digest, "kept_at": kept.KeptAt})
	}
	return map[string]any{"kept": listed}, nil
}

// unpackRequest reads a kept pointer's pieces and the request's bytes, which must match the digest they were kept with.
func unpackRequest(s StateRead, raw []byte) (KeptRequest, []byte, error) {
	var kept KeptRequest
	if err := json.Unmarshal(raw, &kept); err != nil {
		return kept, nil, fmt.Errorf("is no kept request: %w", err)
	}
	if kept.V != 2 {
		return kept, nil, fmt.Errorf("is a kept request of version %d, which this reader does not read (it reads version 2; an older one is past its lifetime within 3 days)", kept.V)
	}
	var packed []byte
	for i := 0; i < kept.Pieces; i++ {
		raw, err := s.Get(RequestChunkKey(kept.Digest, i))
		if err != nil {
			return kept, nil, err
		}
		var piece KeptPiece
		if err := json.Unmarshal(raw, &piece); err != nil || len(raw) == 0 {
			return kept, nil, fmt.Errorf("lacks its piece %d of %d (past its lifetime, or never kept)", i, kept.Pieces)
		}
		part, err := base64.StdEncoding.DecodeString(piece.Gzip)
		if err != nil {
			return kept, nil, fmt.Errorf("holds no base64 in its piece %d: %w", i, err)
		}
		packed = append(packed, part...)
	}
	zr, err := gzip.NewReader(bytes.NewReader(packed))
	if err != nil {
		return kept, nil, fmt.Errorf("holds no gzip: %w", err)
	}
	body, err := io.ReadAll(zr)
	if err != nil {
		return kept, nil, fmt.Errorf("holds no complete gzip: %w", err)
	}
	if sum := sha256.Sum256(body); hex.EncodeToString(sum[:]) != kept.Digest {
		return kept, nil, fmt.Errorf("does not match the digest it was kept with")
	}
	return kept, body, nil
}