/*
DESCRIPTION
  jpeg.go contains code ported from FFmpeg's C implementation of an RTP
  JPEG-compressed Video Depacketizer following RFC 2435. See
  https://ffmpeg.org/doxygen/2.6/rtpdec__jpeg_8c_source.html and
  https://tools.ietf.org/html/rfc2435).

  This code can be used to build JPEG images from an RTP/JPEG stream.

AUTHOR
  Saxon Nelson-Milton <saxon@ausocean.org>

LICENSE
  Copyright (c) 2012 Samuel Pitoiset.

  This file is part of FFmpeg.

  FFmpeg is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.

  FFmpeg is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with FFmpeg; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/

package mjpeg

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

const maxJPEG = 1000000 // 1 MB (arbitrary)

// JPEG marker codes.
const (
	codeSOI  = 0xd8 // Start of image.
	codeDRI  = 0xdd // Define restart interval.
	codeDQT  = 0xdb // Define quantization tables.
	codeDHT  = 0xc4 // Define huffman tables.
	codeSOS  = 0xda // Start of scan.
	codeAPP0 = 0xe0 // TODO: find out what this is.
	codeSOF0 = 0xc0 // Baseline
	codeEOI  = 0xd9 // End of image.
)

var (
	errNoQTable  = errors.New("no quantization table")
	errReservedQ = errors.New("q value is reserved")
)

// Slices used in the creation of huffman tables.
var (
	bitsDCLum = []byte{0, 0, 1, 5, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0}
	bitsDCChr = []byte{0, 0, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0}
	bitsACLum = []byte{0, 0, 2, 1, 3, 3, 2, 4, 3, 5, 5, 4, 4, 0, 0, 1, 0x7d}
	bitsACChr = []byte{0, 0, 2, 1, 2, 4, 4, 3, 4, 7, 5, 4, 4, 0, 1, 2, 0x77}
	valDC     = []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}
	valACLum  = []byte{
		0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12,
		0x21, 0x31, 0x41, 0x06, 0x13, 0x51, 0x61, 0x07,
		0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xa1, 0x08,
		0x23, 0x42, 0xb1, 0xc1, 0x15, 0x52, 0xd1, 0xf0,
		0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0a, 0x16,
		0x17, 0x18, 0x19, 0x1a, 0x25, 0x26, 0x27, 0x28,
		0x29, 0x2a, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39,
		0x3a, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49,
		0x4a, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59,
		0x5a, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69,
		0x6a, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79,
		0x7a, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89,
		0x8a, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98,
		0x99, 0x9a, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7,
		0xa8, 0xa9, 0xaa, 0xb2, 0xb3, 0xb4, 0xb5, 0xb6,
		0xb7, 0xb8, 0xb9, 0xba, 0xc2, 0xc3, 0xc4, 0xc5,
		0xc6, 0xc7, 0xc8, 0xc9, 0xca, 0xd2, 0xd3, 0xd4,
		0xd5, 0xd6, 0xd7, 0xd8, 0xd9, 0xda, 0xe1, 0xe2,
		0xe3, 0xe4, 0xe5, 0xe6, 0xe7, 0xe8, 0xe9, 0xea,
		0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7, 0xf8,
		0xf9, 0xfa,
	}

	valACChr = []byte{
		0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21,
		0x31, 0x06, 0x12, 0x41, 0x51, 0x07, 0x61, 0x71,
		0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91,
		0xa1, 0xb1, 0xc1, 0x09, 0x23, 0x33, 0x52, 0xf0,
		0x15, 0x62, 0x72, 0xd1, 0x0a, 0x16, 0x24, 0x34,
		0xe1, 0x25, 0xf1, 0x17, 0x18, 0x19, 0x1a, 0x26,
		0x27, 0x28, 0x29, 0x2a, 0x35, 0x36, 0x37, 0x38,
		0x39, 0x3a, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48,
		0x49, 0x4a, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58,
		0x59, 0x5a, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68,
		0x69, 0x6a, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78,
		0x79, 0x7a, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
		0x88, 0x89, 0x8a, 0x92, 0x93, 0x94, 0x95, 0x96,
		0x97, 0x98, 0x99, 0x9a, 0xa2, 0xa3, 0xa4, 0xa5,
		0xa6, 0xa7, 0xa8, 0xa9, 0xaa, 0xb2, 0xb3, 0xb4,
		0xb5, 0xb6, 0xb7, 0xb8, 0xb9, 0xba, 0xc2, 0xc3,
		0xc4, 0xc5, 0xc6, 0xc7, 0xc8, 0xc9, 0xca, 0xd2,
		0xd3, 0xd4, 0xd5, 0xd6, 0xd7, 0xd8, 0xd9, 0xda,
		0xe2, 0xe3, 0xe4, 0xe5, 0xe6, 0xe7, 0xe8, 0xe9,
		0xea, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7, 0xf8,
		0xf9, 0xfa,
	}
)

var defaultQuantisers = []byte{
	// Luma table.
	16, 11, 12, 14, 12, 10, 16, 14,
	13, 14, 18, 17, 16, 19, 24, 40,
	26, 24, 22, 22, 24, 49, 35, 37,
	29, 40, 58, 51, 61, 60, 57, 51,
	56, 55, 64, 72, 92, 78, 64, 68,
	87, 69, 55, 56, 80, 109, 81, 87,
	95, 98, 103, 104, 103, 62, 77, 113,
	121, 112, 100, 120, 92, 101, 103, 99,

	/* chroma table */
	17, 18, 18, 24, 21, 24, 47, 26,
	26, 47, 99, 66, 56, 66, 99, 99,
	99, 99, 99, 99, 99, 99, 99, 99,
	99, 99, 99, 99, 99, 99, 99, 99,
	99, 99, 99, 99, 99, 99, 99, 99,
	99, 99, 99, 99, 99, 99, 99, 99,
	99, 99, 99, 99, 99, 99, 99, 99,
	99, 99, 99, 99, 99, 99, 99, 99,
}

// Context describes a RTP/JPEG parsing context that will keep track of the current
// JPEG (held by p), and the state of the quantization tables.
type Context struct {
	qTables    [128][128]byte
	qTablesLen [128]byte
	buf        *putBuffer
	dst        io.Writer
}

// NewContext will return a new Context with destination d.
func NewContext(d io.Writer) *Context {
	return &Context{
		dst: d,
		buf: newPutBuffer(make([]byte, maxJPEG)),
	}
}

// ParsePayload will parse an RTP/JPEG payload and append to current image.
func (c *Context) ParsePayload(p []byte, m bool) error {
	b := newByteStream(p)
	_ = b.get8() // Ignore type-specific flag

	off := b.get24()   // Fragment offset.
	t := b.get8()      // Type.
	q := b.get8()      // Quantization value.
	width := b.get8()  // Picture width.
	height := b.get8() // Picture height.

	var dri int // Restart interval.

	if t&0x40 != 0 {
		dri = b.get16()
		_ = b.get16() // Ignore restart count.
		t &= ^0x40
	}

	if t > 1 {
		panic("unimplemented RTP/JPEG type")
	}

	// Parse quantization table if our offset is 0.
	if off == 0 {
		var qTable []byte
		var qLen int

		if q > 127 {
			_ = b.get8()     // Ignore first byte (reserved for future use).
			prec := b.get8() // The size of coefficients.
			qLen = b.get16() // The length of the quantization table.

			if prec != 0 {
				panic("unsupported precision")
			}

			if qLen > 0 {
				qTable = b.getBuf(qLen)

				if q < 255 {
					if c.qTablesLen[q-128] == 0 && qLen <= 128 {
						copy(c.qTables[q-128][:], qTable)
						c.qTablesLen[q-128] = byte(qLen)
					}
				}
			} else {
				if q == 255 {
					return errNoQTable
				}

				if c.qTablesLen[q-128] == 0 {
					return fmt.Errorf("no quantization tables known for q %d yet", q)
				}

				qTable = c.qTables[q-128][:]
				qLen = int(c.qTablesLen[q-128])
			}
		} else { // q <= 127
			if q == 0 || q > 99 {
				return errReservedQ
			}
			qTable = defaultQTable(q)
			qLen = len(qTable)
		}

		c.buf.reset()

		writeHeader(c.buf, t, width, height, qLen/64, dri, qTable)
	}

	if c.buf.len == 0 {
		// Must have missed start of frame? So ignore and wait for start.
		return nil
	}

	// TODO: check that timestamp is consistent
	// This will need expansion to RTP package to create Timestamp parsing func.

	// TODO: could also check offset with how many bytes we currently have
	// to determine if there are missing frames.

	// Write frame data
	err := b.writeTo(c.buf, b.remaining())
	if err != nil {
		return fmt.Errorf("could not write remaining frame data to output buffer: %w", err)
	}

	if m {
		// End of image marker.
		mark(c.buf, codeEOI)

		_, err = c.buf.writeTo(c.dst)
		if err != nil {
			return fmt.Errorf("could not write JPEG to dst: %w", err)
		}
	}
	return nil
}

// writeHeader writes a JPEG header to the writer w.
func writeHeader(p *putBuffer, _type, width, height, nbqTab, dri int, qtable []byte) {
	width <<= 3
	height <<= 3

	// Indicate start of image.
	mark(p, codeSOI)

	// Write JFIF header.
	mark(p, codeAPP0)
	p.put16(16)
	const jfifLabel = "JFIF\000"
	p.putBuf([]byte(jfifLabel))
	p.put16(0x0201)
	p.put8(0)
	p.put16(1)
	p.put16(1)
	p.put8(0)
	p.put8(0)

	// If we want to define restart interval then write that.
	if dri != 0 {
		mark(p, codeDRI)
		p.put16(4)
		p.put16(uint16(dri))
	}

	// Define quantization tables.
	mark(p, codeDQT)

	// Calculate table size and create slice for table.
	ts := 2 + nbqTab*(1+64)
	p.put16(uint16(ts))

	for i := 0; i < nbqTab; i++ {
		p.put8(uint8(i))
		p.putBuf(qtable[64*i : (64*i)+64])
	}

	// Define huffman table.
	mark(p, codeDHT)
	lenIdx := p.len
	p.put16(0)
	writeHuffman(p, 0, 0, bitsDCLum, valDC)
	writeHuffman(p, 0, 1, bitsDCChr, valDC)
	writeHuffman(p, 1, 0, bitsACLum, valACLum)
	writeHuffman(p, 1, 1, bitsACChr, valACChr)
	p.put16At(uint16(p.len-lenIdx), lenIdx)

	// Start of frame.
	mark(p, codeSOF0)

	// Derive sample type.
	sample := 1
	if _type != 0 {
		sample = 2
	}

	// Derive matrix number.
	mtxNo := 0
	if nbqTab == 2 {
		mtxNo = 1
	}

	p.put16(17)
	p.put8(8)
	p.put16(uint16(height))
	p.put16(uint16(width))
	p.put8(3)
	p.put8(1)
	p.put8(uint8((2 << 4) | sample))
	p.put8(0)
	p.put8(2)
	p.put8(1<<4 | 1)
	p.put8(uint8(mtxNo))
	p.put8(3)
	p.put8(1<<4 | 1)
	p.put8(uint8(mtxNo))

	// Write start of scan.
	mark(p, codeSOS)
	p.put16(12)
	p.put8(3)
	p.put8(1)
	p.put8(0)
	p.put8(2)
	p.put8(17)
	p.put8(3)
	p.put8(17)
	p.put8(0)
	p.put8(63)
	p.put8(0)
}

// mark writes a marker with code c to the putBuffer p.
func mark(p *putBuffer, c byte) {
	p.put8(0xff)
	p.put8(c)
}

// writeHuffman write a JPEG huffman table to w.
func writeHuffman(p *putBuffer, class, id int, bits, values []byte) {
	p.put8(uint8(class<<4 | id))

	var n int
	for i := 1; i <= 16; i++ {
		n += int(bits[i])
	}

	p.putBuf(bits[1:17])
	p.putBuf(values[0:n])
}

// defaultQTable returns a default quantization table.
func defaultQTable(q int) []byte {
	f := clip(q, q, 99)
	const tabLen = 128
	tab := make([]byte, tabLen)

	if q < 50 {
		q = 5000 / f
	} else {
		q = 200 - f*2
	}

	for i := 0; i < tabLen; i++ {
		v := (int(defaultQuantisers[i])*q + 50) / 100
		v = clip(v, 1, 255)
		tab[i] = byte(v)
	}
	return tab
}

// clip clips the value v to the bounds defined by min and max.
func clip(v, min, max int) int {
	if v < min {
		return min
	}

	if v > max {
		return max
	}

	return v
}