packet.go

package dvb

import (
	"fmt"
)

const (
	PESDataIdentifier       = uint8(0x20) // For DVB subtitle streams the data_identifier field shall be coded with the value 0x20
	SubtitleStreamID        = uint8(0x00) // This identifies the subtitle stream in this PES packet. A DVB subtitling stream shall be identified	by the value 0x00
	EndOfPESDataFieldMarker = uint8(0xFF) // An 8-bit field with fixed contents '1111 1111'
)

// Packet represents DVB packet
type Packet struct {
	DataIdentifier              uint8
	SubtitleStreamID            uint8
	PageCompositionSegments     []PageCompositionSegment
	RegionCompositionSegments   []RegionCompositionSegment
	CLUTDefinitionSegments      []CLUTDefinitionSegment
	ObjectDataSegments          []ObjectDataSegment
	DisplayDefinitionSegments   []DisplayDefinitionSegment
	DisparitySignallingSegments []DisparitySignallingSegment
	AlternativeCLUTSegments     []AlternativeCLUTSegment
	EndOfDisplaySegments        []EndOfDisplaySegment
	EndOfPesDataFieldMarker     uint8
}

// Marshal DVB packet into slice of bytes
func (p Packet) Marshal() ([]byte, error) {
	if p.DataIdentifier != PESDataIdentifier {
		return nil, ErrIncorrectPESDataIdentifier
	}

	if p.SubtitleStreamID != SubtitleStreamID {
		return nil, ErrIncorrectSubtitleStreamID
	}

	if p.EndOfPesDataFieldMarker != EndOfPESDataFieldMarker {
		return nil, ErrIncorrectEndOfPESDataFieldMarker
	}

	buf := make([]byte, 0)
	buf = append(buf, PESDataIdentifier, SubtitleStreamID)

	for _, segment := range p.PageCompositionSegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal page composition segment: %w", err)
		}
		buf = append(buf, b...)
	}

	for _, segment := range p.RegionCompositionSegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal region composition segment: %w", err)
		}
		buf = append(buf, b...)
	}

	for _, segment := range p.CLUTDefinitionSegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal CLUT definition segment: %w", err)
		}
		buf = append(buf, b...)
	}

	for _, segment := range p.DisplayDefinitionSegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal display definition segment: %w", err)
		}
		buf = append(buf, b...)
	}

	for _, segment := range p.ObjectDataSegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal segment: %w", err)
		}
		buf = append(buf, b...)
	}

	for _, segment := range p.DisparitySignallingSegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal disparity signalling segment: %w", err)
		}
		buf = append(buf, b...)
	}

	for _, segment := range p.AlternativeCLUTSegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal alternative CLUT segment: %w", err)
		}
		buf = append(buf, b...)
	}

	for _, segment := range p.EndOfDisplaySegments {
		b, err := segment.Marshal()
		if err != nil {
			return buf, fmt.Errorf("failed to marshal end of display segment: %w", err)
		}
		buf = append(buf, b...)
	}

	buf = append(buf, EndOfPESDataFieldMarker)
	return buf, nil
}

func (p *Packet) Unmarshal(buf []byte) error {
	if buf == nil {
		return ErrNilPacket
	}

	if len(buf) < 3 {
		return ErrShortPacket
	}

	p.DataIdentifier = buf[0]
	if p.DataIdentifier != PESDataIdentifier {
		return ErrIncorrectPESDataIdentifier
	}

	p.SubtitleStreamID = buf[1]
	if p.SubtitleStreamID != SubtitleStreamID {
		return ErrIncorrectSubtitleStreamID
	}

	p.EndOfPesDataFieldMarker = buf[len(buf)-1]
	if p.EndOfPesDataFieldMarker != EndOfPESDataFieldMarker {
		return ErrIncorrectEndOfPESDataFieldMarker
	}

	i := newBytesIterator(buf)
	i.Skip(2) // skip 2 header bytes

	for {
		syncByte, err := i.NextByte()
		if err != nil {
			return err
		}
		switch syncByte {
		case segmentSyncByte:
			segmentType, err := i.NextByte()
			if err != nil {
				return err
			}
			switch segmentType {
			case SegmentTypePageComposition:
				_, err := i.NextWordLE() // page id
				if err != nil {
					return err
				}
				segmentLength, err := i.NextWordBE()
				if err != nil {
					return err
				}
				i.Revert(segmentHeaderSize)
				segmentBuf, err := i.NextBytes(segmentHeaderSize + int(segmentLength))
				if err != nil {
					return err
				}
				segment := PageCompositionSegment{}
				if err := segment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal page composition segment: %w", err)
				}
				p.PageCompositionSegments = append(p.PageCompositionSegments, segment)

			case SegmentTypeRegionComposition: // TODO:
				_, err := i.NextWordLE() // page id
				if err != nil {
					return err
				}
				segmentLength, err := i.NextWordBE()
				if err != nil {
					return err
				}
				i.Revert(segmentHeaderSize)
				segmentBuf, err := i.NextBytes(segmentHeaderSize + int(segmentLength))
				if err != nil {
					return err
				}
				segment := RegionCompositionSegment{}
				if err := segment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal region composition segment: %w", err)
				}
				p.RegionCompositionSegments = append(p.RegionCompositionSegments, segment)

			case SegmentTypeCLUTDefinition:
				_, err := i.NextWordLE() // page id
				if err != nil {
					return err
				}
				segmentLength, err := i.NextWordBE()
				if err != nil {
					return err
				}
				i.Revert(segmentHeaderSize)
				segmentBuf, err := i.NextBytes(segmentHeaderSize + int(segmentLength))
				if err != nil {
					return err
				}
				segment := CLUTDefinitionSegment{}
				if err := segment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal CLUT definition segment: %w", err)
				}
				p.CLUTDefinitionSegments = append(p.CLUTDefinitionSegments, segment)

			case SegmentTypeObjectData:
				_, err := i.NextWordLE() // page id
				if err != nil {
					return err
				}
				segmentLength, err := i.NextWordBE()
				if err != nil {
					return err
				}
				i.Revert(segmentHeaderSize)
				segmentBuf, err := i.NextBytes(segmentHeaderSize + int(segmentLength))
				if err != nil {
					return err
				}
				segment := ObjectDataSegment{}
				if err := segment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal object data segment: %w", err)
				}
				p.ObjectDataSegments = append(p.ObjectDataSegments, segment)

			case SegmentTypeDisplayDefinition:
				_, err = i.NextWordBE() // page id
				if err != nil {
					return err
				}
				segmentLength, err := i.NextWordBE()
				if err != nil {
					return err
				}
				i.Revert(segmentHeaderSize)
				segmentBuf, err := i.NextBytes(segmentHeaderSize + int(segmentLength))
				if err != nil {
					return err
				}
				displayDefinitionSegment := DisplayDefinitionSegment{}
				if err := displayDefinitionSegment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal display definition segment: %w", err)
				}
				p.DisplayDefinitionSegments = append(p.DisplayDefinitionSegments, displayDefinitionSegment)

			case SegmentTypeDisparitySignaling:
				_, err = i.NextWordLE() // page id
				if err != nil {
					return err
				}
				segmentLength, err := i.NextWordBE()
				if err != nil {
					return err
				}
				i.Revert(segmentHeaderSize)
				segmentBuf, err := i.NextBytes(segmentHeaderSize + int(segmentLength))
				if err != nil {
					return err
				}
				segment := DisparitySignallingSegment{}
				if err := segment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal disparity signalling segment: %w", err)
				}
				p.DisparitySignallingSegments = append(p.DisparitySignallingSegments, segment)

			case SegmentTypeAlternativeCLUT:
				_, err = i.NextWordLE() // page id
				if err != nil {
					return err
				}
				segmentLength, err := i.NextWordBE()
				if err != nil {
					return err
				}
				i.Revert(segmentHeaderSize)
				segmentBuf, err := i.NextBytes(segmentHeaderSize + int(segmentLength))
				if err != nil {
					return err
				}
				segment := AlternativeCLUTSegment{}
				if err := segment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal alternative CLUT segment: %w", err)
				}
				p.AlternativeCLUTSegments = append(p.AlternativeCLUTSegments, segment)

			case SegmentTypeEndOfDisplay:
				i.Revert(2)
				segmentBuf, err := i.NextBytes(segmentHeaderSize)
				if err != nil {
					return err
				}
				segment := EndOfDisplaySegment{}
				if err := segment.Unmarshal(segmentBuf); err != nil {
					return fmt.Errorf("failed to unmarshal end of display segment: %w", err)
				}
				p.EndOfDisplaySegments = append(p.EndOfDisplaySegments, segment)

			default:
				return ErrIncorrectSegmentType
			}

		case EndOfPESDataFieldMarker:
			return nil

		default:
			return ErrIncorrectSyncByte
		}
	}
}