package colorspace

import (
	"encoding/binary"
	"errors"
	"fmt"
	"strconv"
	"strings"
)

// SubsamplingScheme The subsampling scheme is commonly expressed as a three-part ratio J : A : B (e.g. 4:2:2), that describe the number of luminance and chrominance samples in a conceptual region that is J pixels wide and 2 pixels high.
type SubsamplingScheme struct {
	// J horizontal sampling reference (width of the conceptual region). Usually, 4.
	J byte
	// A number of chrominance samples (Cb, Cr) in the first row of J pixels.
	A byte
	// B number of changes of chrominance samples (Cb, Cr) between first and second row of J pixels. Note that B has to be either zero or equal to A
	B byte
}

type ColorSpace struct {
	Subsampling SubsamplingScheme
	BitDepth    byte
}

func (c ColorSpace) ToInteger() uint32 {
	return binary.BigEndian.Uint32([]byte{c.Subsampling.J, c.Subsampling.A, c.Subsampling.B, c.BitDepth})
}

// ElementPixels returns the number of pixels for a full element
func (c ColorSpace) ElementPixels() int {
	return int(c.Subsampling.J) * 2
}

// ElementSamples returns the number of actual samples (total Y, Cb, Cr) for an encoded element
func (c ColorSpace) ElementSamples() int {
	return int(c.Subsampling.J)*2 + c.ElementChromaSamples()
}

func (c ColorSpace) ElementChromaSamples() int {
	return int(c.Subsampling.A)*2 + int(c.Subsampling.B)*2
}

func (c ColorSpace) PlaneLumaSamples(width, height int) int {
	return width * height
}

func (c ColorSpace) PlaneCrSamples(width, height int) int {
	return int((int64(width) * int64(height) * (int64(c.Subsampling.A) + int64(c.Subsampling.B))) / int64(c.ElementPixels()))
}

func (c ColorSpace) PlaneCbSamples(width, height int) int {
	return int((int64(width) * int64(height) * (int64(c.Subsampling.A) + int64(c.Subsampling.B))) / int64(c.ElementPixels()))
}

func (c ColorSpace) FrameSizePacked(width, height int) (int, error) {
	a1 := int64(width) * int64(height)
	a2 := int64(c.ElementPixels())
	if a1%a2 != 0 {
		return 0, errors.New("not divisible pixels")
	}
	a3 := (a1 / a2) * int64(c.ElementSamples()) * int64(c.BitDepth)
	if a3%8 != 0 {
		return 0, errors.New("not divisible size")
	}

	return int(a3 / 8), nil
}

func (c ColorSpace) FrameSize(width, height int) (int, error) {
	actualBitDepth := int64(c.BitDepth)
	if actualBitDepth&0b111 != 0 {
		actualBitDepth = ((actualBitDepth + 8) >> 3) << 3
	}

	a1 := int64(width) * int64(height)
	a2 := int64(c.ElementPixels())
	if a1%a2 != 0 {
		return 0, errors.New("not divisible pixels")
	}
	a3 := (a1 / a2) * int64(c.ElementSamples()) * actualBitDepth
	if a3%8 != 0 {
		return 0, errors.New("not divisible size")
	}

	return int(a3 / 8), nil
}

func (c ColorSpace) check() error {
	if c.Subsampling.J <= 0 {
		return fmt.Errorf("unsupported J %d", c.Subsampling.J)
	}
	if c.Subsampling.A < 0 && c.Subsampling.A > c.Subsampling.J {
		return fmt.Errorf("unsupported A %d", c.Subsampling.A)
	}
	if c.Subsampling.B != 0 && c.Subsampling.B != c.Subsampling.A {
		return fmt.Errorf("unsupported B %d", c.Subsampling.B)
	}
	if c.BitDepth != 8 && c.BitDepth != 10 && c.BitDepth != 12 && c.BitDepth != 14 && c.BitDepth != 16 {
		return fmt.Errorf("unsupported BitDepth %d", c.BitDepth)
	}

	return nil
}

func NewColorSpaceFromInteger(colorSpace uint32) (ColorSpace, error) {
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, colorSpace)

	space := ColorSpace{
		Subsampling: SubsamplingScheme{
			J: buf[0],
			A: buf[1],
			B: buf[2],
		},
		BitDepth: buf[3],
	}

	return space, space.check()
}

func NewColorSpaceFromString(colorSpace string) (ColorSpace, error) {
	colorSpace = strings.ToLower(colorSpace)
	if colorSpace == "420mpeg2" {
		return NewColorSpaceFromString("420p8")
	}
	splits := strings.Split(colorSpace, "p")
	if len(splits) == 1 { //default 8 bit
		splits = append(splits, "8")
	}

	space := ColorSpace{}
	switch strings.ReplaceAll(splits[0], ":", "") {
	case "420":
		space.Subsampling.J = 4
		space.Subsampling.A = 2
		space.Subsampling.B = 0
	case "422":
		space.Subsampling.J = 4
		space.Subsampling.A = 2
		space.Subsampling.B = 2
	case "444":
		space.Subsampling.J = 4
		space.Subsampling.A = 4
		space.Subsampling.B = 4
	default:
		return space, fmt.Errorf("unsupported Chroma Subsampling %s", splits[0])
	}

	n, err := strconv.Atoi(splits[1])
	if err != nil {
		return space, err
	}
	space.BitDepth = byte(n)

	return space, space.check()
}

func newColorSpaceFromStringInternal(colorSpace string) ColorSpace {
	space, _ := NewColorSpaceFromString(colorSpace)
	return space
}

var (
	Space420    ColorSpace = newColorSpaceFromStringInternal("420")
	Space422    ColorSpace = newColorSpaceFromStringInternal("422")
	Space444    ColorSpace = newColorSpaceFromStringInternal("444")
	Space420P10 ColorSpace = newColorSpaceFromStringInternal("420p10")
	Space422P10 ColorSpace = newColorSpaceFromStringInternal("422p10")
	Space444P10 ColorSpace = newColorSpaceFromStringInternal("444p10")
	Space420P12 ColorSpace = newColorSpaceFromStringInternal("420p12")
	Space422P12 ColorSpace = newColorSpaceFromStringInternal("422p12")
	Space444P12 ColorSpace = newColorSpaceFromStringInternal("444p12")
	Space420P16 ColorSpace = newColorSpaceFromStringInternal("420p16")
	Space422P16 ColorSpace = newColorSpaceFromStringInternal("422p16")
	Space444P16 ColorSpace = newColorSpaceFromStringInternal("444p16")
)