package ingest

import (
	"encoding/binary"
	"fmt"
)

type liveChunk struct {
	FrameType byte
	Payload   []byte
}

type fmp4Assembler struct {
	pending []byte

	initReady bool
	initBuf   []byte

	mediaPrefix []byte
	inFragment  bool
	fragmentBuf []byte
}

func newFMP4Assembler() *fmp4Assembler {
	return &fmp4Assembler{}
}

func (a *fmp4Assembler) Push(data []byte) ([]liveChunk, error) {
	if len(data) == 0 {
		return nil, nil
	}
	a.pending = append(a.pending, data...)

	result := make([]liveChunk, 0, 2)
	for {
		box, boxType, ok, err := a.nextBox()
		if err != nil {
			return nil, err
		}
		if !ok {
			break
		}
		chunks, err := a.consumeBox(boxType, box)
		if err != nil {
			return nil, err
		}
		result = append(result, chunks...)
	}
	return result, nil
}

func (a *fmp4Assembler) nextBox() ([]byte, string, bool, error) {
	if len(a.pending) < 8 {
		return nil, "", false, nil
	}

	headerLen := 8
	size := uint64(binary.BigEndian.Uint32(a.pending[:4]))
	if size == 1 {
		if len(a.pending) < 16 {
			return nil, "", false, nil
		}
		headerLen = 16
		size = binary.BigEndian.Uint64(a.pending[8:16])
	}
	if size == 0 {
		// 0 means box runs to stream end; wait for more data.
		return nil, "", false, nil
	}
	if size < uint64(headerLen) {
		return nil, "", false, fmt.Errorf("invalid MP4 box size %d", size)
	}
	if size > uint64(len(a.pending)) {
		return nil, "", false, nil
	}

	boxLen := int(size)
	boxType := string(a.pending[4:8])
	box := make([]byte, boxLen)
	copy(box, a.pending[:boxLen])
	a.pending = a.pending[boxLen:]
	return box, boxType, true, nil
}

func (a *fmp4Assembler) consumeBox(boxType string, box []byte) ([]liveChunk, error) {
	if !a.initReady {
		a.initBuf = append(a.initBuf, box...)
		if boxType != "moov" {
			return nil, nil
		}

		a.initReady = true
		initPayload := cloneBytes(a.initBuf)
		a.initBuf = nil
		return []liveChunk{{FrameType: liveFrameTypeInit, Payload: initPayload}}, nil
	}

	switch boxType {
	case "styp", "sidx", "prft":
		if a.inFragment {
			a.fragmentBuf = append(a.fragmentBuf, box...)
		} else {
			a.mediaPrefix = append(a.mediaPrefix, box...)
		}
		return nil, nil
	case "moof":
		a.fragmentBuf = a.fragmentBuf[:0]
		if len(a.mediaPrefix) > 0 {
			a.fragmentBuf = append(a.fragmentBuf, a.mediaPrefix...)
			a.mediaPrefix = a.mediaPrefix[:0]
		}
		a.fragmentBuf = append(a.fragmentBuf, box...)
		a.inFragment = true
		return nil, nil
	case "mdat":
		if !a.inFragment {
			return nil, nil
		}
		a.fragmentBuf = append(a.fragmentBuf, box...)
		mediaPayload := cloneBytes(a.fragmentBuf)
		a.fragmentBuf = a.fragmentBuf[:0]
		a.inFragment = false
		return []liveChunk{{FrameType: liveFrameTypeMedia, Payload: mediaPayload}}, nil
	default:
		if a.inFragment {
			a.fragmentBuf = append(a.fragmentBuf, box...)
		}
		return nil, nil
	}
}

func cloneBytes(data []byte) []byte {
	if len(data) == 0 {
		return nil
	}
	out := make([]byte, len(data))
	copy(out, data)
	return out
}