package rdp

// MPPC 解压器实现 RDP 64K 字典压缩（MS-RDPBCGR §3.1.8.4.1）// Logan.Li

const mppcHistorySize = 65536

const (
	mppcBig        = 0x01
	mppcCompressed = 0x20
	mppcReset      = 0x40
	mppcFlushed    = 0x80
)

// MppcDecompressor 维护单连接 MPPC 历史窗口 // Logan.Li
type MppcDecompressor struct {
	history [mppcHistorySize]byte
	offset  int
}

func NewMppcDecompressor() *MppcDecompressor {
	return &MppcDecompressor{}
}

func (d *MppcDecompressor) Decompress(flags byte, data []byte) ([]byte, error) {
	if flags&mppcFlushed != 0 {
		d.history = [mppcHistorySize]byte{}
		d.offset = 0
	} else if flags&mppcReset != 0 {
		d.offset = 0
	}
	if flags&mppcCompressed == 0 {
		for _, b := range data {
			d.history[d.offset] = b
			d.offset = (d.offset + 1) & (mppcHistorySize - 1)
		}
		return data, nil
	}
	output := make([]byte, 0, len(data)*3)
	br := newMppcBitReader(data)
outer:
	for br.bitsLeft >= 9 {
		if br.readBit() == 0 {
			b := byte(br.readBits(8))
			output = append(output, b)
			d.history[d.offset] = b
			d.offset = (d.offset + 1) & (mppcHistorySize - 1)
		} else {
			sel := (br.readBit() << 1) | br.readBit()
			var copyOffset int
			switch sel {
			case 0:
				if br.bitsLeft < 6 {
					break outer
				}
				copyOffset = br.readBits(6)
			case 1:
				if br.bitsLeft < 8 {
					break outer
				}
				copyOffset = 64 + br.readBits(8)
			case 2:
				if br.bitsLeft < 13 {
					break outer
				}
				copyOffset = 320 + br.readBits(13)
			case 3:
				if br.bitsLeft < 16 {
					break outer
				}
				copyOffset = 8192 + br.readBits(16)
			}
			copyLength := 3
			for br.bitsLeft > 0 {
				if br.readBit() == 0 {
					break
				}
				copyLength++
			}
			src := (d.offset - copyOffset - 1 + mppcHistorySize) & (mppcHistorySize - 1)
			for i := 0; i < copyLength; i++ {
				b := d.history[(src+i)&(mppcHistorySize-1)]
				output = append(output, b)
				d.history[d.offset] = b
				d.offset = (d.offset + 1) & (mppcHistorySize - 1)
			}
		}
	}
	return output, nil
}

type mppcBitReader struct {
	data     []byte
	byteIdx  int
	mask     byte
	bitsLeft int
}

func newMppcBitReader(data []byte) *mppcBitReader {
	return &mppcBitReader{data: data, mask: 0x80, bitsLeft: len(data) * 8}
}

func (r *mppcBitReader) readBit() int {
	if r.bitsLeft == 0 {
		return 0
	}
	r.bitsLeft--
	bit := 0
	if r.data[r.byteIdx]&r.mask != 0 {
		bit = 1
	}
	r.mask >>= 1
	if r.mask == 0 {
		r.mask = 0x80
		r.byteIdx++
	}
	return bit
}

func (r *mppcBitReader) readBits(n int) int {
	result := 0
	for i := 0; i < n; i++ {
		result = (result << 1) | r.readBit()
	}
	return result
}