package tpm

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math"

	"github.com/google/go-tpm/legacy/tpm2"

	"go.step.sm/crypto/internal/utils/convert"
)

type ShortRandomReadError struct {
	Requested int
	Generated int
}

func (s ShortRandomReadError) Error() string {
	return fmt.Sprintf("generated %d random bytes instead of the requested %d", s.Generated, s.Requested)
}

// GenerateRandom returns `size` number of random bytes generated by the TPM.
func (t *TPM) GenerateRandom(ctx context.Context, size uint16) (random []byte, err error) {
	if err = t.open(goTPMCall(ctx)); err != nil {
		return nil, fmt.Errorf("failed opening TPM: %w", err)
	}
	defer closeTPM(ctx, t, &err)

	return t.generateRandom(ctx, size)
}

func (t *TPM) generateRandom(_ context.Context, size uint16) (random []byte, err error) {
	random, err = tpm2.GetRandom(t.rwc, size)
	if err != nil {
		return nil, fmt.Errorf("failed generating random data: %w", err)
	}

	if len(random) != int(size) {
		return nil, ShortRandomReadError{Requested: int(size), Generated: len(random)}
	}

	return
}

type generator struct {
	t         *TPM
	readError error
}

func (t *TPM) RandomReader() (io.Reader, error) {
	return &generator{
		t: t,
	}, nil
}

func (g *generator) Read(p []byte) (n int, err error) {
	if g.readError != nil {
		errMsg := g.readError.Error() // multiple wrapped errors not (yet) allowed
		return 0, fmt.Errorf("failed generating random bytes in previous call to Read: %s: %w", errMsg, io.EOF)
	}
	if len(p) > math.MaxUint16 {
		p = p[:math.MaxUint16]
	}

	ctx := context.Background()
	if err = g.t.open(goTPMCall(ctx)); err != nil {
		return 0, fmt.Errorf("failed opening TPM: %w", err)
	}
	defer closeTPM(ctx, g.t, &err)

	var result []byte
	requestedLength := len(p)
	singleRequestLength := convert.MustUint16(requestedLength)
	for len(result) < requestedLength {
		if r, err := g.t.generateRandom(ctx, singleRequestLength); err == nil {
			result = append(result, r...)
		} else {
			var s ShortRandomReadError
			if errors.As(err, &s) && s.Generated > 0 {
				// adjust number of bytes to request if at least some data was read and continue loop
				singleRequestLength = convert.MustUint16(s.Generated)
				result = append(result, r...)
			} else {
				g.readError = err // store the error to be returned for future calls to Read
				n = copy(p, result)
				return n, nil // return the result recorded so far and no error
			}
		}
	}

	n = copy(p, result)
	return
}

var _ io.Reader = (*generator)(nil)
