// GoToSocial
// Copyright (C) GoToSocial Authors admin@gotosocial.org
// SPDX-License-Identifier: AGPL-3.0-or-later
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program 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 Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package media

import (
	"context"
	"encoding/binary"
	"image/jpeg"
	"io"
	"strings"

	"code.superseriousbusiness.org/gopkg/log"
	"code.superseriousbusiness.org/gotosocial/internal/gtserror"
	"codeberg.org/gruf/go-byteutil"
)

const (
	// image magic header bytes.
	magicJPEG = "\xff\xd8\xff"
)

// probe will first attempt to probe the file at path using native Go code
// (for performance), but falls back to using ffprobe to retrieve media details.
func probe(ctx context.Context, filepath string) (*result, error) {

	// Open input file at given path.
	file, err := openRead(filepath)
	if err != nil {
		return nil, gtserror.Newf("error opening file %s: %w", filepath, err)
	}

	// Close on return.
	defer file.Close()

	// Byte buf to check for
	// file header magic bytes.
	buf := make([]byte, 3)

	// Read file header into buffer.
	_, err = io.ReadFull(file, buf)
	if err != nil {
		return nil, gtserror.Newf("error reading file %s: %w", filepath, err)
	}

	switch {
	// Attempt to probe JPEG types
	// separately, to save calls into
	// WebAssembly for a common image.
	case string(buf[:len(magicJPEG)]) == magicJPEG:
		log.Debug(ctx, "probing jpeg")
		return probeJPEG(file)

	default:
		// Close BEFORE
		// pass to ffprobe.
		_ = file.Close()

		// For everything else, fall back
		// to calling ffprobe on input file.
		log.Debug(ctx, "ffprobing file")
		return ffprobe(ctx, filepath)
	}
}

// probeJPEG decodes the given file as JPEG and determines
// image details from the decoded JPEG using native Go code.
func probeJPEG(file fileReader) (*result, error) {

	// Attempt to decode JPEG, adding back hdr magic.
	cfg, err := jpeg.DecodeConfig(io.MultiReader(
		strings.NewReader(magicJPEG),
		file,
	))
	if err != nil {
		err := gtserror.Newf("error decoding file %s: %w", file.Name(), err)
		return nil, withDetails(err, codecDetails)
	}

	// Jump back to file start.
	_, err = file.Seek(0, 0)
	if err != nil {
		return nil, gtserror.Newf("error seeking in file %s: %w", file.Name(), err)
	}

	// Read orientation data from EXIF.
	orientation := readOrientation(file)

	// Setup result as if
	// ffprobe'd resulting in
	// JPEG file container.
	var res result
	res.format = "image2"

	// Set image orientation data.
	res.orientation = orientation

	// Extract image details.
	res.video = []videoStream{{
		stream: stream{codec: "mjpeg"},
		width:  cfg.Width,
		height: cfg.Height,

		// setting a pixel color format
		// doesn't matter for JPEG, as we
		// don't bother even using it.
		pixfmt: "",
	}}

	return &res, nil
}

// readOrientation reads orientation EXIF
// data (if it even exists) from image file.
//
// copied from github.com/disintegration/imaging
// but modified to optimize discard operations.
func readOrientation(r fileReader) int {
	const (
		markerAPP1     = 0xffe1
		exifHeader     = 0x45786966
		byteOrderBE    = 0x4d4d
		byteOrderLE    = 0x4949
		orientationTag = 0x0112
	)

	// Setup a read buffer.
	var buf byteutil.Buffer
	buf.B = make([]byte, 0, 64)

	// discard simply reads into buf.
	discard := func(n int) error {
		buf.Guarantee(n)
		_, err := io.ReadFull(r, buf.B[:n])
		return err
	}

	// readUint16 reads uint16 bytes into buffer then parses.
	readUint16 := func(b binary.ByteOrder) (uint16, error) {
		_, err := io.ReadFull(r, buf.B[:2])
		return b.Uint16(buf.B[:2]), err
	}

	// readUint32 reads uint32 bytes into buffer then parses.
	readUint32 := func(b binary.ByteOrder) (uint32, error) {
		_, err := io.ReadFull(r, buf.B[:4])
		return b.Uint32(buf.B[:4]), err
	}

	// Skip past JPEG SOI marker.
	if err := discard(2); err != nil {
		return orientationUnspecified
	}

	// Find JPEG
	// APP1 marker.
	for {
		marker, err := readUint16(binary.BigEndian)
		if err != nil {
			return orientationUnspecified
		}

		size, err := readUint16(binary.BigEndian)
		if err != nil {
			return orientationUnspecified
		}

		if marker>>8 != 0xff {
			return orientationUnspecified // Invalid JPEG marker.
		}

		if marker == markerAPP1 {
			break
		}

		if size < 2 {
			return orientationUnspecified // Invalid block size.
		}

		if err := discard(int(size - 2)); err != nil {
			return orientationUnspecified
		}
	}

	// Check if EXIF header is present.
	header, err := readUint32(binary.BigEndian)
	if err != nil {
		return orientationUnspecified
	}

	if header != exifHeader {
		return orientationUnspecified
	}

	if err := discard(2); err != nil {
		return orientationUnspecified
	}

	// Read byte order info.
	byteOrderTag, err := readUint16(binary.BigEndian)
	if err != nil {
		return orientationUnspecified
	}

	var byteOrder binary.ByteOrder
	switch byteOrderTag {
	case byteOrderBE:
		byteOrder = binary.BigEndian
	case byteOrderLE:
		byteOrder = binary.LittleEndian
	default:
		return orientationUnspecified // Invalid byte order flag.
	}

	if err := discard(2); err != nil {
		return orientationUnspecified
	}

	// Skip the EXIF offset.
	offset, err := readUint32(byteOrder)
	if err != nil {
		return orientationUnspecified
	}

	if offset < 8 {
		return orientationUnspecified // Invalid offset value.
	}

	if err := discard(int(offset - 8)); err != nil {
		return orientationUnspecified
	}

	// Read the number of tags.
	numTags, err := readUint16(byteOrder)
	if err != nil {
		return orientationUnspecified
	}

	// Find the orientation tag.
	for i := 0; i < int(numTags); i++ {
		tag, err := readUint16(byteOrder)
		if err != nil {
			return orientationUnspecified
		}

		if tag != orientationTag {
			if err := discard(10); err != nil {
				return orientationUnspecified
			}
			continue
		}

		if err := discard(6); err != nil {
			return orientationUnspecified
		}

		val, err := readUint16(byteOrder)
		if err != nil {
			return orientationUnspecified
		}

		if val < 1 || val > 8 {
			return orientationUnspecified // Invalid tag value.
		}

		return int(val)
	}

	// Missing orientation tag.
	return orientationUnspecified
}
