package builder

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"hash/crc32"
	"io"
	"os"
)

// xzBlocksMax and xzStreamsMax are the most blocks and the most concatenated streams an xz file may hold for
// the build seat to read it.
//
// **Why a block count is limited at all.** The xz reader (github.com/therootcompany/xz v1.0.1) keeps one
// filter decoder for every block that names a delta or BCJ filter, and frees them only when a stream ends:
// its block counters are reset per stream, not per block. A stream of many tiny blocks therefore grows the
// heap with its length: an `xz --block-size=1 --delta` file of 6 MB kept about 67 MiB live, which comes to
// about 5.6 GiB at the 512 MiB download cap (review 68 on mesh-controller pull request 264, hq issue 345).
//
// **Why 1,024.** At about 340 bytes of decoder a block, 1,024 blocks keep at most about 0.35 MiB live. A
// real release is far below it: Betterbird's archive is 1 stream with 1 block, and multi-threaded xz cuts
// blocks of three times the dictionary (24 MiB at the default preset), so even 512 MiB unpacked in blocks
// of 512 KiB stays within it. Streams are limited to 16 because a release is one stream; the decoders are
// reused from one stream to the next, so the stream limit only bounds the work of this walk.
const (
	xzBlocksMax  = 1024
	xzStreamsMax = 16
)

// xzWithinLimits walks an xz file's framing, without decompressing anything, and refuses one holding more
// blocks or streams than the build seat reads.
//
// **The framing, not the index.** Each stream's index at its end lists its blocks, but the reader compares
// the index with the blocks only after it has decoded them all, so a file whose index lies would still be
// decoded block by block. This walk follows the same framing the reader follows: stream header, each block
// header, the block's LZMA2 chunks (each chunk header carries its own compressed size, so a chunk is skipped
// without being decoded), its padding and check, then the index and the footer. It counts the blocks the
// reader would decode, and it also refuses an index whose count of blocks is not the count walked.
func xzWithinLimits(file string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()
	_, _, err = xzFraming(bufio.NewReader(f))
	if errors.Is(err, errXzTooManyBlocks) || errors.Is(err, errXzTooManyStreams) {
		return err
	}
	if err != nil {
		return fmt.Errorf("it is not an xz file this build seat reads: %v", err)
	}
	return nil
}

var (
	errXzTooManyBlocks = fmt.Errorf("it holds more than %d xz blocks, the most this build seat reads: "+
		"its xz reader keeps a filter decoder for every block (hq issue 345)", xzBlocksMax)
	errXzTooManyStreams = fmt.Errorf("it holds more than %d concatenated xz streams, the most this build "+
		"seat reads", xzStreamsMax)
)

// xzFraming counts the streams and blocks of an xz file, checking its framing as it goes. It stops with
// errXzTooManyBlocks or errXzTooManyStreams at the first block or stream over the limit.
func xzFraming(r *bufio.Reader) (streams, blocks int, err error) {
	full := func(n int) ([]byte, error) {
		b := make([]byte, n)
		if _, err := io.ReadFull(r, b); err != nil {
			if err == io.EOF {
				err = io.ErrUnexpectedEOF
			}
			return nil, err
		}
		return b, nil
	}
	skip := func(n int) error {
		if got, err := r.Discard(n); err != nil || got != n {
			return io.ErrUnexpectedEOF
		}
		return nil
	}
	for {
		header, err := full(12)
		if err != nil {
			return streams, blocks, fmt.Errorf("stream %d: header: %v", streams+1, err)
		}
		if !bytes.Equal(header[:6], []byte{0xfd, '7', 'z', 'X', 'Z', 0x00}) {
			return streams, blocks, fmt.Errorf("stream %d: no xz stream header", streams+1)
		}
		if header[6] != 0 || header[7]&0xf0 != 0 ||
			crc32.ChecksumIEEE(header[6:8]) != binary.LittleEndian.Uint32(header[8:12]) {
			return streams, blocks, fmt.Errorf("stream %d: a corrupt stream header", streams+1)
		}
		if streams++; streams > xzStreamsMax {
			return streams, blocks, errXzTooManyStreams
		}
		checkSize := 0
		if check := int(header[7]); check != 0 {
			checkSize = 4 << ((check - 1) / 3)
		}
		inStream := 0
		for {
			first, err := r.ReadByte()
			if err != nil {
				return streams, blocks, fmt.Errorf("stream %d: no index: %v", streams, io.ErrUnexpectedEOF)
			}
			if first == 0 {
				break // the index indicator
			}
			rest, err := full(int(first)*4 + 3)
			if err != nil {
				return streams, blocks, fmt.Errorf("stream %d, block %d: header: %v", streams, inStream+1, err)
			}
			if err := xzBlockHeader(append([]byte{first}, rest...)); err != nil {
				return streams, blocks, fmt.Errorf("stream %d, block %d: %v", streams, inStream+1, err)
			}
			inStream++
			if blocks++; blocks > xzBlocksMax {
				return streams, blocks, errXzTooManyBlocks
			}
			// The block's LZMA2 chunks, up to the end marker.
			size := 0
			for {
				control, err := r.ReadByte()
				if err != nil {
					return streams, blocks, fmt.Errorf("stream %d, block %d: %v", streams, inStream, io.ErrUnexpectedEOF)
				}
				size++
				if control == 0 {
					break
				}
				switch {
				case control == 1 || control == 2: // stored: a 16-bit size less one, then the bytes
					b, err := full(2)
					if err != nil {
						return streams, blocks, err
					}
					n := int(binary.BigEndian.Uint16(b)) + 1
					if err := skip(n); err != nil {
						return streams, blocks, err
					}
					size += 2 + n
				case control >= 0x80: // LZMA: unpacked and packed sizes, the properties byte on a reset
					b, err := full(4)
					if err != nil {
						return streams, blocks, err
					}
					n := int(binary.BigEndian.Uint16(b[2:])) + 1
					size += 4
					if control >= 0xc0 {
						if err := skip(1); err != nil {
							return streams, blocks, err
						}
						size++
					}
					if err := skip(n); err != nil {
						return streams, blocks, err
					}
					size += n
				default:
					return streams, blocks, fmt.Errorf("stream %d, block %d: an LZMA2 chunk of kind %#x",
						streams, inStream, control)
				}
			}
			if pad := (4 - size%4) % 4; pad > 0 {
				b, err := full(pad)
				if err != nil || !bytes.Equal(b, make([]byte, pad)) {
					return streams, blocks, fmt.Errorf("stream %d, block %d: corrupt block padding", streams, inStream)
				}
			}
			if err := skip(checkSize); err != nil {
				return streams, blocks, err
			}
		}
		// The index: its count of records, two numbers each, padding, CRC32.
		used := 1
		count, n, err := xzVLI(r)
		if err != nil {
			return streams, blocks, fmt.Errorf("stream %d: index: %v", streams, err)
		}
		used += n
		if count != uint64(inStream) {
			return streams, blocks, fmt.Errorf("stream %d: its index lists %d blocks, but it holds %d",
				streams, count, inStream)
		}
		for i := 0; i < 2*inStream; i++ {
			_, n, err := xzVLI(r)
			if err != nil {
				return streams, blocks, fmt.Errorf("stream %d: index: %v", streams, err)
			}
			used += n
		}
		if err := skip((4-used%4)%4 + 4); err != nil {
			return streams, blocks, fmt.Errorf("stream %d: index: %v", streams, err)
		}
		footer, err := full(12)
		if err != nil || footer[10] != 'Y' || footer[11] != 'Z' {
			return streams, blocks, fmt.Errorf("stream %d: no stream footer", streams)
		}
		// Stream padding: zero bytes, a multiple of four, then the end or another stream.
		zeros := 0
		for {
			b, err := r.ReadByte()
			if err == io.EOF {
				if zeros%4 != 0 {
					return streams, blocks, errors.New("stream padding not a multiple of four bytes")
				}
				return streams, blocks, nil
			}
			if err != nil {
				return streams, blocks, err
			}
			if b != 0 {
				_ = r.UnreadByte()
				break
			}
			zeros++
		}
		if zeros%4 != 0 {
			return streams, blocks, errors.New("stream padding not a multiple of four bytes")
		}
	}
}

// xzBlockHeader checks one block header: its CRC32, its flags, and that its last filter is LZMA2, the only
// one whose chunks this walk can skip (and the only last filter the reader decodes).
func xzBlockHeader(h []byte) error {
	size := len(h)
	if crc32.ChecksumIEEE(h[:size-4]) != binary.LittleEndian.Uint32(h[size-4:]) {
		return errors.New("a corrupt block header")
	}
	flags := h[1]
	if flags&0x3c != 0 {
		return errors.New("a block header with reserved flags")
	}
	body := bytes.NewReader(h[2 : size-4])
	if flags&0x40 != 0 {
		if _, _, err := xzVLI(body); err != nil {
			return err
		}
	}
	if flags&0x80 != 0 {
		if _, _, err := xzVLI(body); err != nil {
			return err
		}
	}
	var last uint64
	for i := 0; i <= int(flags&3); i++ {
		id, _, err := xzVLI(body)
		if err != nil {
			return err
		}
		props, _, err := xzVLI(body)
		if err != nil {
			return err
		}
		if props > uint64(body.Len()) {
			return errors.New("a corrupt block header")
		}
		if _, err := body.Seek(int64(props), io.SeekCurrent); err != nil {
			return err
		}
		last = id
	}
	if last != 0x21 {
		return fmt.Errorf("a block whose last filter is %#x, not LZMA2", last)
	}
	return nil
}

// xzVLI reads one of xz's variable-length integers (at most 9 bytes, 7 bits each), and how many bytes it took.
func xzVLI(r io.ByteReader) (uint64, int, error) {
	var v uint64
	for i := 0; i < 9; i++ {
		b, err := r.ReadByte()
		if err != nil {
			return 0, i, io.ErrUnexpectedEOF
		}
		v |= uint64(b&0x7f) << (7 * i)
		if b&0x80 == 0 {
			if i > 0 && b == 0 {
				return 0, i + 1, errors.New("a number not in its shortest form")
			}
			return v, i + 1, nil
		}
	}
	return 0, 9, errors.New("a number longer than nine bytes")
}