package gateway

import (
	"bytes"
	"encoding/binary"
)

const (
	apolloVideoMaximumDataShards = 255
	apolloVideoMaximumBlocks     = 4
	apolloVideoShardPayloadSize  = apolloVideoRawPacketSize - apolloRTPHeaderSize - 4 - apolloVideoNVHeaderSize
)

type apolloVideoShard struct {
	frame       uint32
	block       uint8
	lastBlock   uint8
	dataPackets int
	parity      int
	index       int
	sequence    uint16
	streamIndex uint32
	flags       byte
	payload     []byte
}

type apolloVideoFECBlock struct {
	dataPackets int
	parity      int
	firstSeq    uint16
	streamBase  uint32
	haveBase    bool
	shards      [][]byte
	received    []bool
	count       int
	complete    bool
}

type apolloVideoAssembler struct {
	haveFrame bool
	frame     uint32
	lastBlock uint8
	blocks    [apolloVideoMaximumBlocks]*apolloVideoFECBlock
}

func (a *apolloVideoAssembler) Add(shard apolloVideoShard) ([]byte, error) {
	if len(shard.payload) != apolloVideoShardPayloadSize || shard.dataPackets < 1 || shard.dataPackets > apolloVideoMaximumDataShards || shard.parity < 0 || shard.dataPackets+shard.parity > 255 || shard.block > shard.lastBlock || shard.lastBlock >= apolloVideoMaximumBlocks || shard.index >= shard.dataPackets+shard.parity {
		return nil, errApolloMedia
	}
	if !a.haveFrame || apolloFrameNewer(shard.frame, a.frame) {
		*a = apolloVideoAssembler{haveFrame: true, frame: shard.frame, lastBlock: shard.lastBlock}
	} else if shard.frame != a.frame {
		return nil, nil
	}
	if shard.lastBlock != a.lastBlock {
		return nil, errApolloMedia
	}
	block := a.blocks[shard.block]
	if block == nil {
		block = &apolloVideoFECBlock{
			dataPackets: shard.dataPackets,
			parity:      shard.parity,
			shards:      make([][]byte, shard.dataPackets+shard.parity),
			received:    make([]bool, shard.dataPackets+shard.parity),
		}
		a.blocks[shard.block] = block
	} else if block.dataPackets != shard.dataPackets || block.parity != shard.parity {
		return nil, errApolloMedia
	}
	if shard.index < shard.dataPackets {
		if shard.index == 0 && shard.flags&0x04 == 0 || shard.index == shard.dataPackets-1 && shard.flags&0x02 == 0 || shard.flags&^byte(0x07) != 0 {
			return nil, errApolloMedia
		}
		base := shard.sequence - uint16(shard.index)
		streamBase := shard.streamIndex - uint32(shard.index)
		if !block.haveBase {
			block.firstSeq, block.streamBase, block.haveBase = base, streamBase, true
		} else if block.firstSeq != base || block.streamBase != streamBase {
			return nil, errApolloMedia
		}
	}
	if block.received[shard.index] {
		if !bytes.Equal(block.shards[shard.index], shard.payload) {
			return nil, errApolloMedia
		}
		return nil, nil
	}
	block.shards[shard.index] = append([]byte(nil), shard.payload...)
	block.received[shard.index] = true
	block.count++
	if block.count < block.dataPackets {
		return nil, nil
	}
	if !block.complete {
		if err := reconstructApolloVideoBlock(block); err != nil {
			return nil, err
		}
		block.complete = true
	}
	for index := uint8(0); index <= a.lastBlock; index++ {
		if a.blocks[index] == nil || !a.blocks[index].complete {
			return nil, nil
		}
	}
	capacity := 0
	for blockIndex := uint8(0); blockIndex <= a.lastBlock; blockIndex++ {
		capacity += a.blocks[blockIndex].dataPackets * apolloVideoShardPayloadSize
	}
	frame := make([]byte, 0, capacity)
	for blockIndex := uint8(0); blockIndex <= a.lastBlock; blockIndex++ {
		for shardIndex := 0; shardIndex < a.blocks[blockIndex].dataPackets; shardIndex++ {
			frame = append(frame, a.blocks[blockIndex].shards[shardIndex]...)
		}
	}
	if len(frame) < 8 || frame[0] != 0x01 {
		return nil, errApolloMedia
	}
	lastPayloadLength := int(binary.LittleEndian.Uint16(frame[4:6]))
	end := len(frame) - apolloVideoShardPayloadSize + lastPayloadLength
	if lastPayloadLength == 0 || lastPayloadLength > apolloVideoShardPayloadSize || end <= 8 || end > len(frame) {
		return nil, errApolloMedia
	}
	*a = apolloVideoAssembler{}
	return append([]byte(nil), frame[8:end]...), nil
}

func apolloFrameNewer(first, second uint32) bool {
	return int32(first-second) > 0
}

func reconstructApolloVideoBlock(block *apolloVideoFECBlock) error {
	if block == nil || block.count < block.dataPackets {
		return errApolloMedia
	}
	missing := false
	for index := 0; index < block.dataPackets; index++ {
		if !block.received[index] {
			missing = true
			block.shards[index] = make([]byte, apolloVideoShardPayloadSize)
		}
	}
	if !missing {
		return nil
	}
	selectedRows := make([][]byte, 0, block.dataPackets)
	selectedShards := make([][]byte, 0, block.dataPackets)
	for index, received := range block.received {
		if !received {
			continue
		}
		selectedRows = append(selectedRows, apolloVideoFECRow(index, block.dataPackets, block.parity))
		selectedShards = append(selectedShards, block.shards[index])
		if len(selectedRows) == block.dataPackets {
			break
		}
	}
	if len(selectedRows) != block.dataPackets {
		return errApolloMedia
	}
	inverse, ok := apolloGFInvert(selectedRows)
	if !ok {
		return errApolloMedia
	}
	for index := 0; index < block.dataPackets; index++ {
		if block.received[index] {
			continue
		}
		for source, coefficient := range inverse[index] {
			apolloGFAXPY(block.shards[index], selectedShards[source], coefficient)
		}
		block.received[index] = true
	}
	return nil
}

func apolloVideoFECRow(index, dataPackets, parity int) []byte {
	row := make([]byte, dataPackets)
	if index < dataPackets {
		row[index] = 1
		return row
	}
	parityIndex := index - dataPackets
	for dataIndex := range row {
		row[dataIndex] = apolloGFInverse(byte((parity + dataIndex) ^ parityIndex))
	}
	return row
}

func apolloGFInvert(matrix [][]byte) ([][]byte, bool) {
	size := len(matrix)
	if size == 0 {
		return nil, false
	}
	work := make([][]byte, size)
	inverse := make([][]byte, size)
	for row := range matrix {
		if len(matrix[row]) != size {
			return nil, false
		}
		work[row] = append([]byte(nil), matrix[row]...)
		inverse[row] = make([]byte, size)
		inverse[row][row] = 1
	}
	for column := 0; column < size; column++ {
		pivot := column
		for pivot < size && work[pivot][column] == 0 {
			pivot++
		}
		if pivot == size {
			return nil, false
		}
		work[column], work[pivot] = work[pivot], work[column]
		inverse[column], inverse[pivot] = inverse[pivot], inverse[column]
		factor := apolloGFInverse(work[column][column])
		for index := column; index < size; index++ {
			work[column][index] = apolloGFMultiply(work[column][index], factor)
		}
		for index := range inverse[column] {
			inverse[column][index] = apolloGFMultiply(inverse[column][index], factor)
		}
		for row := 0; row < size; row++ {
			if row == column || work[row][column] == 0 {
				continue
			}
			factor = work[row][column]
			for index := column; index < size; index++ {
				work[row][index] ^= apolloGFMultiply(work[column][index], factor)
			}
			for index := range inverse[row] {
				inverse[row][index] ^= apolloGFMultiply(inverse[column][index], factor)
			}
		}
	}
	return inverse, true
}

func apolloGFAXPY(destination, source []byte, coefficient byte) {
	if coefficient == 0 {
		return
	}
	for index := range destination {
		destination[index] ^= apolloGFMultiply(source[index], coefficient)
	}
}

func apolloGFInverse(value byte) byte {
	if value == 0 {
		return 0
	}
	return apolloGFPow(value, 254)
}

func apolloGFPow(value byte, exponent uint8) byte {
	result := byte(1)
	for exponent != 0 {
		if exponent&1 != 0 {
			result = apolloGFMultiply(result, value)
		}
		value = apolloGFMultiply(value, value)
		exponent >>= 1
	}
	return result
}

func apolloGFMultiply(first, second byte) byte {
	var product byte
	for second != 0 {
		if second&1 != 0 {
			product ^= first
		}
		high := first & 0x80
		first <<= 1
		if high != 0 {
			first ^= 0x1d
		}
		second >>= 1
	}
	return product
}