package ndi

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"image"
	"image/color"
	"image/color/palette"
	"io"
	"log"
)

var (
	ErrNotNDI         = errors.New("ndi: not an ndi file")
	ErrUnexpectedData = errors.New("ndi: unexpected data")

	NDI1MagicBytes = []byte("ndi1")

	Logger = log.New(log.Writer(), "ndi debug: ", 0)
)

const (
	LUTColorStride = 3
	HeaderStride   = 12
	LUTPadding     = 4
)

func init() {
	image.RegisterFormat("ndi", string(NDI1MagicBytes), Decode, DecodeConfig)
}

func Decode(r io.Reader) (o image.Image, err error) {
	buf := bytes.Buffer{}
	_, err = io.Copy(&buf, r)
	if err != nil {
		return o, fmt.Errorf("ndi: decode copy failed")
	}

	cfgBuf := bytes.NewBuffer(buf.Bytes())
	cfg, err := DecodeConfig(cfgBuf)
	if err != nil {
		return o, fmt.Errorf("ndi: decode config failed: %w", err)
	}

	cfgBuf = bytes.NewBuffer(buf.Bytes())
	palette, err := DecodePalette(cfgBuf)
	if err != nil {
		return o, fmt.Errorf("ndi: decode palette failed: %w", err)
	}

	bounds := image.Rect(0, 0, cfg.Width, cfg.Height)
	img := image.NewPaletted(bounds, palette)

	// skip until end of LUT
	buf.Next(HeaderStride + (len(palette) * LUTColorStride) + LUTPadding)

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			index, err := buf.ReadByte()
			if err != nil {
				// skip because we don't mind bad pixels
				continue
			}

			img.SetColorIndex(x, y, uint8(index)%uint8(len(palette)))
		}
	}

	return img, nil
}

func DecodeConfig(r io.Reader) (cfg image.Config, err error) {
	imageBuf := bytes.Buffer{}
	readBytes, err := io.Copy(&imageBuf, r)
	if err != nil {
		return cfg, fmt.Errorf("ndi: decode config copy failed")
	}

	header := imageBuf.Next(HeaderStride)

	magicBytes := header[0:4]
	widthBytes := header[4:6]
	lutSizeByte := header[6]

	if !bytes.Equal(magicBytes, NDI1MagicBytes) {
		return cfg, ErrNotNDI
	}

	width := binary.BigEndian.Uint16(widthBytes)
	lutByteWidth := int(lutSizeByte) * LUTColorStride

	imagePixelCount := readBytes - ((int64(lutByteWidth) + LUTPadding) + HeaderStride)
	height := imagePixelCount / int64(width)

	cfg.ColorModel = color.RGBAModel
	cfg.Width = int(width)
	cfg.Height = int(height)

	return cfg, nil
}

func DecodePalette(r io.Reader) (p color.Palette, err error) {
	imageBuf := bytes.Buffer{}
	_, err = io.Copy(&imageBuf, r)
	if err != nil {
		return
	}

	header := imageBuf.Next(HeaderStride)
	lutCount := int(header[6])

	tableSize := lutCount * LUTColorStride
	tableBytes := imageBuf.Next(tableSize)

	p = make(color.Palette, lutCount)

	for index := 0; index < lutCount; index++ {
		baseIndex := index * LUTColorStride

		p[index] = color.RGBA{
			R: tableBytes[baseIndex+0],
			G: tableBytes[baseIndex+1],
			B: tableBytes[baseIndex+2],
			A: 0xFF,
		}
	}

	return
}

func Encode(w io.Writer, img image.Image, cfg *Options) error {
	if cfg == nil {
		cfg = &Options{
			Palette: palette.WebSafe,
		}
	}

	_, err := w.Write(NDI1MagicBytes)
	if err != nil {
		return err
	}

	bounds := img.Bounds()
	width := bounds.Max.X - bounds.Min.X
	widthBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(widthBytes, uint16(width))

	Logger.Println("width ", width)
	Logger.Println("width bytes ", widthBytes)

	_, err = w.Write(widthBytes)
	if err != nil {
		return err
	}

	err = encodeLUT(w, cfg)
	if err != nil {
		return err
	}

	w.Write([]byte("AWAW"))

	err = encodePixels(w, img, cfg)
	if err != nil {
		return err
	}

	return nil
}

func encodeLUT(w io.Writer, cfg *Options) error {
	if len(cfg.Palette) == 0 {
		cfg.Palette = palette.WebSafe
	}

	_, err := w.Write([]byte{byte(len(cfg.Palette))})
	if err != nil {
		return err
	}

	w.Write([]byte("AWAWA"))

	for _, color := range cfg.Palette {
		r, g, b, _ := color.RGBA()
		_, err = w.Write([]byte{
			byte(r / 256),
			byte(g / 256),
			byte(b / 256),
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func encodePixels(w io.Writer, img image.Image, cfg *Options) error {
	bounds := img.Bounds()

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			color := img.At(x, y)
			index := uint8(cfg.Palette.Index(color))
			_, err := w.Write([]byte{index})
			if err != nil {
				return err
			}
		}
	}

	return nil
}

type Options struct {
	Palette color.Palette
}