package imaging

import (
	"bytes"
	"encoding/binary"
	"hash/crc32"
)

const (
	pngChunkHeaderSize = 8
	pngChunkCRCSize    = 4
	pngIHDR            = "IHDR"
	pngIEND            = "IEND"
	pngEXIF            = "eXIf"
)

var pngSignature = []byte("\x89PNG\r\n\x1a\n")

var pngMetadataChunks = map[string]bool{
	pngEXIF: true,
	"tEXt":  true,
	"zTXt":  true,
	"iTXt":  true,
	"tIME":  true,
}

func stripPNG(data []byte) ([]byte, error) {
	if !bytes.HasPrefix(data, pngSignature) {
		return nil, errMalformedContainer
	}

	out := bytes.Clone(pngSignature)
	orientation := uint16(uprightOrientation)
	insertAt := 0

	for pos := len(pngSignature); pos < len(data); {
		available := len(data) - pos - pngChunkHeaderSize - pngChunkCRCSize
		if available < 0 {
			return nil, errMalformedContainer
		}
		size := binary.BigEndian.Uint32(data[pos:])
		if uint64(size) > uint64(available) {
			return nil, errMalformedContainer
		}
		kind := string(data[pos+4 : pos+pngChunkHeaderSize])
		payload := data[pos+pngChunkHeaderSize : pos+pngChunkHeaderSize+int(size)]
		end := pos + pngChunkHeaderSize + len(payload) + pngChunkCRCSize

		if pngMetadataChunks[kind] {
			if kind == pngEXIF {
				if found, ok := exifOrientation(payload); ok {
					orientation = found
				}
			}
		} else {
			out = append(out, data[pos:end]...)
			if kind == pngIHDR {
				insertAt = len(out)
			}
		}
		pos = end

		if kind == pngIEND {
			return withPNGOrientation(out, insertAt, orientation)
		}
	}
	return nil, errMalformedContainer
}

func withPNGOrientation(out []byte, insertAt int, orientation uint16) ([]byte, error) {
	if !needsOrientation(orientation) {
		return out, nil
	}
	if insertAt == 0 {
		return nil, errMalformedContainer
	}
	chunk := pngChunk(pngEXIF, orientationOnlyTIFF(orientation))
	return append(out[:insertAt:insertAt], append(chunk, out[insertAt:]...)...), nil
}

func pngChunk(kind string, payload []byte) []byte {
	chunk := binary.BigEndian.AppendUint32(nil, uint32(len(payload)))
	chunk = append(chunk, kind...)
	chunk = append(chunk, payload...)
	return binary.BigEndian.AppendUint32(chunk, crc32.ChecksumIEEE(chunk[4:]))
}