package attachment

import (
	"bytes"
	"encoding/binary"
	"image"
	"io"
)

// readExifOrientation extracts the EXIF Orientation tag (0x0112) from a JPEG
// blob. It returns 0 when the tag is absent or anything fails to parse —
// callers treat every value except 2..8 as "no transform".
func readExifOrientation(jpegBytes []byte) uint16 {
	reader := bytes.NewReader(jpegBytes)
	// SOI signature.
	if _, err := reader.Seek(2, io.SeekStart); err != nil || len(jpegBytes) < 2 ||
		jpegBytes[0] != 0xFF || jpegBytes[1] != 0xD8 {
		return 0
	}
	for {
		marker, payload, err := readJPEGSegment(reader)
		if err != nil {
			return 0
		}
		if marker != 0xE1 { // APP1
			continue
		}
		if len(payload) < 8 || string(payload[0:6]) != "Exif\x00\x00" {
			continue
		}
		return orientationFromTIFF(payload[6:])
	}
}

// readJPEGSegment returns the next length-bearing metadata segment's marker
// and payload. JPEG encoders may place COM, DQT, or other segments before
// APP1, so callers must be able to skip all of them rather than only APPn. It
// stops at the scan header: EXIF only ever lives in the leading metadata.
func readJPEGSegment(reader *bytes.Reader) (byte, []byte, error) {
	for {
		prefix, err := reader.ReadByte()
		if err != nil {
			return 0, nil, err
		}
		if prefix != 0xFF {
			continue
		}
		marker, err := reader.ReadByte()
		if err != nil {
			return 0, nil, err
		}
		if marker == 0xFF || marker == 0x00 {
			continue
		}
		// SOS begins entropy-coded image data and EOI ends the image; metadata
		// following either marker is not part of the leading EXIF area.
		if marker == 0xDA || marker == 0xD9 {
			return 0, nil, io.EOF
		}
		// TEM, restart markers, and a repeated SOI have no length field.
		if marker == 0x01 || (marker >= 0xD0 && marker <= 0xD8) {
			continue
		}
		var length [2]byte
		if _, err := io.ReadFull(reader, length[:]); err != nil {
			return 0, nil, err
		}
		size := int(length[0])<<8 | int(length[1])
		if size < 2 {
			return 0, nil, io.ErrUnexpectedEOF
		}
		payload := make([]byte, size-2)
		if _, err := io.ReadFull(reader, payload); err != nil {
			return 0, nil, err
		}
		return marker, payload, nil
	}
}

func orientationFromTIFF(tiff []byte) uint16 {
	if len(tiff) < 8 {
		return 0
	}
	var byteOrder binary.ByteOrder
	switch string(tiff[0:2]) {
	case "II":
		byteOrder = binary.LittleEndian
	case "MM":
		byteOrder = binary.BigEndian
	default:
		return 0
	}
	if byteOrder.Uint16(tiff[2:4]) != 42 { // TIFF magic
		return 0
	}
	ifdOffset := int(byteOrder.Uint32(tiff[4:8]))
	if ifdOffset < 8 || ifdOffset+2 > len(tiff) {
		return 0
	}
	entryCount := int(byteOrder.Uint16(tiff[ifdOffset : ifdOffset+2]))
	for i := 0; i < entryCount; i++ {
		entryOffset := ifdOffset + 2 + i*12
		if entryOffset+12 > len(tiff) {
			return 0
		}
		entry := tiff[entryOffset : entryOffset+12]
		if byteOrder.Uint16(entry[0:2]) != 0x0112 { // Orientation
			continue
		}
		if byteOrder.Uint16(entry[2:4]) != 3 { // SHORT
			return 0
		}
		if byteOrder.Uint32(entry[4:8]) != 1 { // exactly one value
			return 0
		}
		return byteOrder.Uint16(entry[8:10])
	}
	return 0
}

// applyExifOrientation returns a copy of src baked to its upright display
// orientation, so re-encoded derivatives lose the EXIF tag without losing the
// orientation itself. Orientation 0/1 (and any unexpected value) returns src.
func applyExifOrientation(src image.Image, orientation uint16) image.Image {
	switch orientation {
	case 2, 3, 4, 5, 6, 7, 8:
	default:
		return src
	}
	bounds := src.Bounds()
	w, h := bounds.Dx(), bounds.Dy()
	if orientation >= 5 {
		w, h = h, w
	}
	dst := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			srcX, srcY := exifSourceCoordinates(int(orientation), x, y, bounds.Dx(), bounds.Dy())
			dst.Set(x, y, src.At(bounds.Min.X+srcX, bounds.Min.Y+srcY))
		}
	}
	return dst
}

// exifSourceCoordinates maps an upright destination pixel back to its stored
// source coordinates for orientations 2..8. Values 5..8 swap the axes, so x
// ranges over srcH and y over srcW there.
func exifSourceCoordinates(orientation, x, y, srcW, srcH int) (int, int) {
	switch orientation {
	case 2: // mirror horizontal
		return srcW - 1 - x, y
	case 3: // rotate 180
		return srcW - 1 - x, srcH - 1 - y
	case 4: // mirror vertical
		return x, srcH - 1 - y
	case 5: // mirror horizontal + rotate 270 CW = transpose
		return y, x
	case 6: // rotate 90 CW
		return y, srcH - 1 - x
	case 7: // mirror horizontal + rotate 90 CW = anti-transpose
		return srcW - 1 - y, srcH - 1 - x
	case 8: // rotate 270 CW
		return srcW - 1 - y, x
	default:
		return x, y
	}
}