// Copyright The Monogon Project Authors.
// SPDX-License-Identifier: Apache-2.0

package blockdev

import (
	"errors"
	"fmt"
	"io"
)

var ErrNotBlockDevice = errors.New("not a block device")

// BlockDev represents a generic block device made up of equally-sized blocks.
// All offsets and intervals are expressed in bytes and must be aligned to
// BlockSize and are recommended to be aligned to OptimalBlockSize if feasible.
// Unless stated otherwise, intervals are inclusive-exclusive, i.e. the
// start byte is included but the end byte is not.
type BlockDev interface {
	io.ReaderAt
	io.WriterAt

	// BlockCount returns the number of blocks on the block device or -1 if it
	// is an image with an undefined size.
	BlockCount() int64

	// BlockSize returns the block size of the block device in bytes. This must
	// be a power of two and is commonly (but not always) either 512 or 4096.
	BlockSize() int64

	// OptimalBlockSize returns the optimal block size in bytes for aligning
	// to as well as issuing I/O. IO operations with block sizes below this
	// one might incur read-write overhead. This is the larger of the physical
	// block size and a device-reported value if available.
	OptimalBlockSize() int64

	// Discard discards a continuous set of blocks. Discarding means the
	// underlying device gets notified that the data in these blocks is no
	// longer needed. This can improve performance of the device device (as it
	// no longer needs to preserve the unused data) as well as bulk erase
	// operations. This command is advisory and not all implementations support
	// it. The contents of discarded blocks are implementation-defined.
	Discard(startByte int64, endByte int64) error

	// Zero zeroes a continouous set of blocks. On certain implementations this
	// can be significantly faster than just calling Write with zeroes.
	Zero(startByte, endByte int64) error

	// Sync commits the current contents to stable storage.
	Sync() error
}

func NewRWS(b BlockDev) *ReadWriteSeeker {
	return &ReadWriteSeeker{b: b}
}

// ReadWriteSeeker provides an adapter implementing ReadWriteSeeker on top of
// a blockdev.
type ReadWriteSeeker struct {
	b       BlockDev
	currPos int64
}

func (s *ReadWriteSeeker) Read(p []byte) (n int, err error) {
	n, err = s.b.ReadAt(p, s.currPos)
	s.currPos += int64(n)
	return
}

func (s *ReadWriteSeeker) Write(p []byte) (n int, err error) {
	n, err = s.b.WriteAt(p, s.currPos)
	s.currPos += int64(n)
	return
}

func (s *ReadWriteSeeker) Seek(offset int64, whence int) (int64, error) {
	switch whence {
	default:
		return 0, errors.New("Seek: invalid whence")
	case io.SeekStart:
	case io.SeekCurrent:
		offset += s.currPos
	case io.SeekEnd:
		offset += s.b.BlockCount() * s.b.BlockSize()
	}
	if offset < 0 {
		return 0, errors.New("Seek: invalid offset")
	}
	s.currPos = offset
	return s.currPos, nil
}

var ErrOutOfBounds = errors.New("write out of bounds")

// NewSection returns a new Section, implementing BlockDev over that subset
// of blocks. The interval is inclusive-exclusive.
func NewSection(b BlockDev, startBlock, endBlock int64) (*Section, error) {
	if startBlock < 0 {
		return nil, fmt.Errorf("invalid range: startBlock (%d) negative", startBlock)
	}
	if startBlock > endBlock {
		return nil, fmt.Errorf("invalid range: startBlock (%d) bigger than endBlock (%d)", startBlock, endBlock)
	}
	if endBlock > b.BlockCount() {
		return nil, fmt.Errorf("endBlock (%d) out of range (%d)", endBlock, b.BlockCount())
	}
	return &Section{
		b:          b,
		startBlock: startBlock,
		endBlock:   endBlock,
	}, nil
}

// Section implements BlockDev on a slice of another BlockDev given a startBlock
// and endBlock.
type Section struct {
	b                    BlockDev
	startBlock, endBlock int64
}

func (s *Section) ReadAt(p []byte, off int64) (n int, err error) {
	if off < 0 {
		return 0, errors.New("blockdev.Section.ReadAt: negative offset")
	}
	bOff := off + (s.startBlock * s.b.BlockSize())
	bytesToEnd := (s.endBlock * s.b.BlockSize()) - bOff
	if bytesToEnd < 0 {
		return 0, io.EOF
	}
	if bytesToEnd < int64(len(p)) {
		n, err := s.b.ReadAt(p[:bytesToEnd], bOff)
		if err == nil {
			err = io.EOF
		}
		return n, err
	}
	return s.b.ReadAt(p, bOff)
}

func (s *Section) WriteAt(p []byte, off int64) (n int, err error) {
	bOff := off + (s.startBlock * s.b.BlockSize())
	bytesToEnd := (s.endBlock * s.b.BlockSize()) - bOff
	if off < 0 || bytesToEnd < 0 {
		return 0, ErrOutOfBounds
	}
	if bytesToEnd < int64(len(p)) {
		n, err := s.b.WriteAt(p[:bytesToEnd], bOff)
		if err != nil {
			// If an error happened, prioritize that error
			return n, err
		}
		// Otherwise, return ErrOutOfBounds as even short writes must return an
		// error.
		return n, ErrOutOfBounds
	}
	return s.b.WriteAt(p, bOff)
}

func (s *Section) BlockCount() int64 {
	return s.endBlock - s.startBlock
}

func (s *Section) BlockSize() int64 {
	return s.b.BlockSize()
}

func (s *Section) OptimalBlockSize() int64 {
	return s.b.OptimalBlockSize()
}

func (s *Section) Discard(startByte, endByte int64) error {
	if err := validAlignedRange(s, startByte, endByte); err != nil {
		return err
	}
	offset := s.startBlock * s.b.BlockSize()
	return s.b.Discard(offset+startByte, offset+endByte)
}

func (s *Section) Zero(startByte, endByte int64) error {
	if err := validAlignedRange(s, startByte, endByte); err != nil {
		return err
	}
	offset := s.startBlock * s.b.BlockSize()
	return s.b.Zero(offset+startByte, offset+endByte)
}

func (s *Section) Sync() error {
	return s.b.Sync()
}

func validAlignedRange(b BlockDev, startByte, endByte int64) error {
	if startByte < 0 {
		return fmt.Errorf("invalid range: startByte (%d) negative", startByte)
	}
	if startByte > endByte {
		return fmt.Errorf("invalid range: startByte (%d) bigger than endByte (%d)", startByte, endByte)
	}
	devLen := b.BlockCount() * b.BlockSize()
	if endByte > devLen {
		return fmt.Errorf("endByte (%d) out of range (%d)", endByte, devLen)
	}
	if startByte%b.BlockSize() != 0 {
		return fmt.Errorf("startByte (%d) needs to be aligned to block size (%d)", startByte, b.BlockSize())
	}
	if endByte%b.BlockSize() != 0 {
		return fmt.Errorf("endByte (%d) needs to be aligned to block size (%d)", endByte, b.BlockSize())
	}
	return nil
}

// GenericZero implements software-based zeroing. This can be used to implement
// Zero when no acceleration is available or desired.
func GenericZero(b BlockDev, startByte, endByte int64) error {
	if err := validAlignedRange(b, startByte, endByte); err != nil {
		return err
	}
	// Choose buffer size close to 16MiB or the range to be zeroed, whatever
	// is smaller.
	bufSizeTarget := int64(16 * 1024 * 1024)
	if endByte-startByte < bufSizeTarget {
		bufSizeTarget = endByte - startByte
	}
	bufSize := (bufSizeTarget / b.BlockSize()) * b.BlockSize()
	buf := make([]byte, bufSize)
	for i := startByte; i < endByte; i += bufSize {
		if endByte-i < bufSize {
			buf = buf[:endByte-i]
		}
		if _, err := b.WriteAt(buf, i); err != nil {
			return fmt.Errorf("while writing zeroes: %w", err)
		}
	}
	return nil
}