package document

import (
	"Scanyonero/unit"
	"encoding/binary"
	"fmt"
	"io"

	"neilpa.me/go-jfif"
)

type jfifTagUnits byte

const (
	jfifTagUnitsNoUnits jfifTagUnits = iota
	jfifTagUnitsDotsPerInch
	jfifTagUnitsDotsPerCM
)

type jfifTag struct {
	Version  uint16
	Units    jfifTagUnits
	XDensity uint16
	YDensity uint16
	// Omit thumbnail width, height and pixel data.
}

func (j *jfifTag) UnmarshalBinary(data []byte) error {
	if len(data) < 7 {
		return fmt.Errorf("JFIF tag length (%d) is smaller than expected (%d)", len(data), 7)
	}

	j.Version = binary.BigEndian.Uint16(data[0:2])
	j.Units = jfifTagUnits(data[2])
	j.XDensity = binary.BigEndian.Uint16(data[3:5])
	j.YDensity = binary.BigEndian.Uint16(data[5:7])

	return nil
}

// Density returns the number of pixels per length unit.
func (j *jfifTag) Density() (x, y unit.Density) {
	switch j.Units {
	case jfifTagUnitsDotsPerInch:
		return unit.PerInch(j.XDensity), unit.PerInch(j.YDensity)
	case jfifTagUnitsDotsPerCM:
		return unit.PerMillimeter(j.XDensity) / 10, unit.PerMillimeter(j.YDensity) / 10
	}

	return nil, nil
}

func decodeJFIF(r io.Reader) (jfifTag, error) {
	segments, err := jfif.DecodeSegments(r)
	if err != nil {
		return jfifTag{}, fmt.Errorf("failed to decode JPEG segments: %w", err)
	}

	for _, segment := range segments {
		sig, data, _ := segment.AppPayload()
		switch sig {
		case jfif.SigJFIF:
			tag := jfifTag{}
			if err := tag.UnmarshalBinary(data); err != nil {
				return jfifTag{}, fmt.Errorf("failed to unmarshal tag data: %w", err)
			}
			return tag, nil
		}
	}

	return jfifTag{}, fmt.Errorf("couldn't find any JFIF tag")
}