package psbt

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
)

// Errors the reader refuses.
var (
	// ErrNotPSBT is a blob that is not a BIP174 file.
	ErrNotPSBT = errors.New("psbt: not a psbt")
	// ErrTruncated is a record whose length runs past the end of the file.
	ErrTruncated = errors.New("psbt: truncated")
)

// Preimage is one SHA-256 preimage record: the hash it opens and the preimage.
type Preimage struct {
	// SecretHash is the hash the record is keyed by.
	SecretHash [32]byte
	// Preimage is the preimage.
	Preimage []byte
}

// Input is what a reader got out of one input map.
type Input struct {
	// Amount and ScriptPubKey come from the witness UTXO record, if it is there.
	Amount       uint64
	ScriptPubKey []byte
	// WitnessScript is the P2WSH script, if the map carries one.
	WitnessScript []byte
	// Preimages are the preimage records.
	Preimages []Preimage
}

// File is a parsed PSBT.
type File struct {
	// Unsigned is the global unsigned transaction, as serialised.
	Unsigned []byte
	// Inputs are the input maps, in order.
	Inputs []Input
	// Outputs is how many output maps were present.
	Outputs int
}

// Read parses a PSBT.
//
// It is the other direction of Write, and it checks what it can check: magic, that every length
// fits, and that a preimage record's key really is the hash of its value. A record whose
// meaning it does not know is skipped rather than guessed at — a PSBT that carries something
// this wallet did not write is not an error, it is another signer's business.
func Read(raw []byte) (*File, error) {
	if !bytes.HasPrefix(raw, magic) {
		return nil, ErrNotPSBT
	}

	reader := &cursor{raw: raw, at: len(magic)}
	file := &File{}

	// The global map.
	for {
		key, value, err := reader.record()

		if err != nil {
			return nil, err
		}

		if key == nil {
			break
		}

		switch key[0] {
		case globalUnsignedTx:
			file.Unsigned = value
		}
	}

	// BIP174 does not count the maps, so the unsigned transaction does: it says how many input
	// maps and how many output maps follow.
	inputs, outputs, err := counts(file.Unsigned)

	if err != nil {
		return nil, err
	}

	for index := 0; index < inputs; index++ {
		input := Input{}
		emptied := false

		for {
			if !reader.more() {
				return nil, fmt.Errorf("%w: an unclosed input map", ErrTruncated)
			}

			key, value, err := reader.record()

			if err != nil {
				return nil, err
			}

			if key == nil {
				emptied = true

				break
			}

			switch key[0] {
			case inputWitnessUTXO:
				if len(value) < 9 {
					return nil, fmt.Errorf("%w: a witness utxo of %d bytes", ErrTruncated, len(value))
				}

				input.Amount = binary.LittleEndian.Uint64(value[0:8])

				length := int(value[8])

				if len(value) < 9+length {
					return nil, fmt.Errorf("%w: a witness utxo script", ErrTruncated)
				}

				input.ScriptPubKey = value[9 : 9+length]
			case inputWitnessScript:
				input.WitnessScript = value
			case inputSHA256Preimage:
				if len(key) != 33 {
					return nil, fmt.Errorf("%w: a preimage key of %d bytes", ErrTruncated, len(key))
				}

				record := Preimage{Preimage: value}
				copy(record.SecretHash[:], key[1:])

				if !bytes.Equal(sha256Of(value), record.SecretHash[:]) {
					return nil, fmt.Errorf("%w: a preimage does not open %x", ErrTruncated, record.SecretHash)
				}

				input.Preimages = append(input.Preimages, record)
			}
		}

		if !emptied {
			break
		}

		file.Inputs = append(file.Inputs, input)
	}

	for index := 0; index < outputs; index++ {
		for {
			if !reader.more() {
				return nil, fmt.Errorf("%w: an unclosed output map", ErrTruncated)
			}

			key, _, err := reader.record()

			if err != nil {
				return nil, err
			}

			if key == nil {
				file.Outputs++

				break
			}
		}
	}

	if reader.more() {
		return nil, fmt.Errorf("%w: %d trailing bytes", ErrTruncated, len(reader.raw)-reader.at)
	}

	return file, nil
}

// cursor walks the file.
type cursor struct {
	raw []byte
	at  int
}

// more reports whether there are bytes left.
func (c *cursor) more() bool { return c.at < len(c.raw) }

// record reads one key and value, or a nil key at a map's end.
func (c *cursor) record() ([]byte, []byte, error) {
	if !c.more() {
		return nil, nil, nil
	}

	if c.raw[c.at] == 0x00 {
		c.at++

		return nil, nil, nil
	}

	length, err := c.size()

	if err != nil {
		return nil, nil, err
	}

	key, err := c.take(length)

	if err != nil {
		return nil, nil, err
	}

	length, err = c.size()

	if err != nil {
		return nil, nil, err
	}

	value, err := c.take(length)

	if err != nil {
		return nil, nil, err
	}

	return key, value, nil
}

// size reads a compact size.
func (c *cursor) size() (int, error) {
	if !c.more() {
		return 0, ErrTruncated
	}

	first := c.raw[c.at]
	c.at++

	switch first {
	case 0xfd:
		raw, err := c.take(2)

		if err != nil {
			return 0, err
		}

		return int(binary.LittleEndian.Uint16(raw)), nil
	case 0xfe:
		raw, err := c.take(4)

		if err != nil {
			return 0, err
		}

		return int(binary.LittleEndian.Uint32(raw)), nil
	}

	return int(first), nil
}

// take reads n bytes.
func (c *cursor) take(n int) ([]byte, error) {
	if n < 0 || c.at+n > len(c.raw) {
		return nil, fmt.Errorf("%w: %d bytes at %d", ErrTruncated, n, c.at)
	}

	out := c.raw[c.at : c.at+n]
	c.at += n

	return out, nil
}

// counts walks the unsigned transaction and reports how many inputs and outputs it has, which
// is how many maps must follow.
func counts(unsigned []byte) (int, int, error) {
	cursor := &cursor{raw: unsigned}

	if _, err := cursor.take(4); err != nil {
		return 0, 0, err
	}

	inputs, err := cursor.size()

	if err != nil {
		return 0, 0, err
	}

	for index := 0; index < inputs; index++ {
		if _, err := cursor.take(36); err != nil { // outpoint
			return 0, 0, err
		}

		length, err := cursor.size()

		if err != nil {
			return 0, 0, err
		}

		if _, err := cursor.take(length); err != nil { // scriptSig
			return 0, 0, err
		}

		if _, err := cursor.take(4); err != nil { // sequence
			return 0, 0, err
		}
	}

	outputs, err := cursor.size()

	if err != nil {
		return 0, 0, err
	}

	for index := 0; index < outputs; index++ {
		if _, err := cursor.take(8); err != nil { // value
			return 0, 0, err
		}

		length, err := cursor.size()

		if err != nil {
			return 0, 0, err
		}

		if _, err := cursor.take(length); err != nil { // script
			return 0, 0, err
		}
	}

	return inputs, outputs, nil
}