package imaging

import (
	"bytes"
	"encoding/binary"
)

const (
	jpegMarkerPrefix = 0xFF
	jpegSOI          = 0xD8
	jpegEOI          = 0xD9
	jpegSOS          = 0xDA
	jpegTEM          = 0x01
	jpegRST0         = 0xD0
	jpegRST7         = 0xD7
	jpegAPP0         = 0xE0
	jpegAPP1         = 0xE1
	jpegAPP2         = 0xE2
	jpegAPP14        = 0xEE
	jpegAPP15        = 0xEF
	jpegCOM          = 0xFE
	jpegStuffedZero  = 0x00
)

var iccProfileSignature = []byte("ICC_PROFILE\x00")

type jpegSegment struct {
	marker  byte
	payload []byte
}

func stripJPEG(data []byte) ([]byte, error) {
	if len(data) < 4 || data[0] != jpegMarkerPrefix || data[1] != jpegSOI {
		return nil, errMalformedContainer
	}

	out := []byte{jpegMarkerPrefix, jpegSOI}
	insertAt := len(out)
	orientation := uint16(uprightOrientation)
	first := true

	for pos := 2; ; {
		marker, next, err := readJPEGMarker(data, pos)
		if err != nil {
			return nil, err
		}
		pos = next

		if marker == jpegEOI {
			out = append(out, jpegMarkerPrefix, jpegEOI)
			return withOrientationSegment(out, insertAt, orientation), nil
		}
		if isStandaloneJPEGMarker(marker) {
			out = append(out, jpegMarkerPrefix, marker)
			continue
		}

		segment, next, err := readJPEGSegment(data, pos, marker)
		if err != nil {
			return nil, err
		}
		pos = next

		if isJPEGMetadata(segment) {
			if found, ok := jpegExifOrientation(segment); ok {
				orientation = found
			}
		} else {
			out = appendJPEGSegment(out, segment)
			if first && segment.marker == jpegAPP0 {
				insertAt = len(out)
			}
		}
		first = false

		if marker == jpegSOS {
			end := jpegScanEnd(data, pos)
			out = append(out, data[pos:end]...)
			if end == len(data) {
				return withOrientationSegment(out, insertAt, orientation), nil
			}
			pos = end
		}
	}
}

func readJPEGMarker(data []byte, pos int) (byte, int, error) {
	if pos >= len(data) || data[pos] != jpegMarkerPrefix {
		return 0, 0, errMalformedContainer
	}
	for pos < len(data) && data[pos] == jpegMarkerPrefix {
		pos++
	}
	if pos >= len(data) {
		return 0, 0, errMalformedContainer
	}
	return data[pos], pos + 1, nil
}

func readJPEGSegment(data []byte, pos int, marker byte) (jpegSegment, int, error) {
	if pos+2 > len(data) {
		return jpegSegment{}, 0, errMalformedContainer
	}
	length := int(binary.BigEndian.Uint16(data[pos:]))
	if length < 2 || pos+length > len(data) {
		return jpegSegment{}, 0, errMalformedContainer
	}
	return jpegSegment{marker: marker, payload: data[pos+2 : pos+length]}, pos + length, nil
}

func jpegScanEnd(data []byte, pos int) int {
	for pos < len(data) {
		if data[pos] != jpegMarkerPrefix || pos+1 >= len(data) {
			pos++
			continue
		}
		next := data[pos+1]
		if next == jpegStuffedZero || isRestartMarker(next) {
			pos += 2
			continue
		}
		return pos
	}
	return len(data)
}

func isStandaloneJPEGMarker(marker byte) bool {
	return marker == jpegTEM || isRestartMarker(marker)
}

func isRestartMarker(marker byte) bool {
	return marker >= jpegRST0 && marker <= jpegRST7
}

func isJPEGMetadata(segment jpegSegment) bool {
	if segment.marker == jpegCOM {
		return true
	}
	if segment.marker == jpegAPP2 {
		return !bytes.HasPrefix(segment.payload, iccProfileSignature)
	}
	if segment.marker == jpegAPP0 || segment.marker == jpegAPP14 {
		return false
	}
	return segment.marker >= jpegAPP1 && segment.marker <= jpegAPP15
}

func jpegExifOrientation(segment jpegSegment) (uint16, bool) {
	if segment.marker != jpegAPP1 || !bytes.HasPrefix(segment.payload, exifHeader) {
		return 0, false
	}
	return exifOrientation(segment.payload)
}

func appendJPEGSegment(out []byte, segment jpegSegment) []byte {
	out = append(out, jpegMarkerPrefix, segment.marker)
	out = binary.BigEndian.AppendUint16(out, uint16(len(segment.payload)+2))
	return append(out, segment.payload...)
}

func withOrientationSegment(out []byte, insertAt int, orientation uint16) []byte {
	if !needsOrientation(orientation) {
		return out
	}
	payload := append(bytes.Clone(exifHeader), orientationOnlyTIFF(orientation)...)
	segment := appendJPEGSegment(nil, jpegSegment{marker: jpegAPP1, payload: payload})
	return append(out[:insertAt:insertAt], append(segment, out[insertAt:]...)...)
}