package x264

/*
#cgo pkg-config: x264
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <x264.h>
*/
import "C"
import (
	"errors"
	"fmt"
	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
	"git.gammaspectra.live/S.O.N.G/Ignite/y4m"
	"io"
	"runtime"
	"strconv"
	"sync/atomic"
	"unsafe"
)

type Encoder struct {
	w          io.Writer
	cleaned    atomic.Bool
	params     C.x264_param_t
	pictureIn  *C.x264_picture_t
	pictureOut *C.x264_picture_t
	h          *C.x264_t
}

func NewEncoder(w io.Writer, stream *y4m.Stream, settings map[string]any) (*Encoder, error) {
	e := &Encoder{
		w: w,
	}

	preset := C.CString(getSettingString(settings, "preset", "medium"))
	defer C.free(unsafe.Pointer(preset))

	if C.x264_param_default_preset(&e.params, preset, nil) < 0 {
		return nil, errors.New("error setting preset")
	}

	defaultProfile := "high"

	switch true {
	case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 4 && stream.ColorSpace().Subsampling.B == 4:
		e.params.i_csp = C.X264_CSP_I444
		defaultProfile = "high444"
	case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 2:
		e.params.i_csp = C.X264_CSP_I422
		defaultProfile = "high422"
	case stream.ColorSpace().Subsampling.J == 4 && stream.ColorSpace().Subsampling.A == 2 && stream.ColorSpace().Subsampling.B == 0:
		e.params.i_csp = C.X264_CSP_I420
		defaultProfile = "high"
	default:
		return nil, errors.New("unsupported input chroma subsampling")

	}

	profile := C.CString(getSettingString(settings, "profile", defaultProfile))
	defer C.free(unsafe.Pointer(profile))

	if stream.ColorSpace().BitDepth > 8 {
		e.params.i_csp |= C.X264_CSP_HIGH_DEPTH
		if defaultProfile == "high" {
			defaultProfile = "high10"
		}
	}
	e.params.i_bitdepth = C.int(stream.ColorSpace().BitDepth)

	width, height := stream.Resolution()
	e.params.i_width = C.int(width)
	e.params.i_height = C.int(height)
	e.params.b_vfr_input = 0
	e.params.b_repeat_headers = 1
	e.params.b_annexb = 1
	e.params.i_fps_num = C.uint32_t(stream.FrameRate().Numerator)
	e.params.i_fps_den = C.uint32_t(stream.FrameRate().Denominator)

	if stream.ColorRange() {
		e.params.vui.b_fullrange = 1
	} else {
		e.params.vui.b_fullrange = 0
	}

	if C.x264_param_apply_profile(&e.params, profile) < 0 {
		return nil, errors.New("error setting profile")
	}

	for k, v := range settings {
		if err := func() error {
			var strVal *C.char
			if val, ok := v.(string); ok {
				strVal = C.CString(val)
			} else if val, ok := v.(int); ok {
				strVal = C.CString(strconv.Itoa(val))
			} else if val, ok := v.(int64); ok {
				strVal = C.CString(strconv.Itoa(int(val)))
			}

			if strVal != nil {
				defer C.free(unsafe.Pointer(strVal))
			} else {
				return fmt.Errorf("could not get parameter %s", k)
			}
			strKey := C.CString(k)
			defer C.free(unsafe.Pointer(strKey))
			if ret := C.x264_param_parse(&e.params, strKey, strVal); ret != 0 {
				if ret == C.X264_PARAM_BAD_NAME {
					return fmt.Errorf("bad parameter name %s", k)
				} else if ret == C.X264_PARAM_BAD_VALUE {
					return fmt.Errorf("bad parameter value %s for %s", C.GoString(strVal), k)
				} else {
					return fmt.Errorf("error setting parameter %s", k)
				}
			}

			return nil
		}(); err != nil {
			return nil, err
		}
	}

	runtime.SetFinalizer(e, func(encoder *Encoder) {
		encoder.Close()
	})

	if e.pictureIn = (*C.x264_picture_t)(C.malloc(C.size_t(unsafe.Sizeof(C.x264_picture_t{})))); e.pictureIn == nil {
		return nil, errors.New("error allocating memory")
	}
	if e.pictureOut = (*C.x264_picture_t)(C.malloc(C.size_t(unsafe.Sizeof(C.x264_picture_t{})))); e.pictureOut == nil {
		return nil, errors.New("error allocating memory")
	}
	if C.x264_picture_alloc(e.pictureIn, e.params.i_csp, e.params.i_width, e.params.i_height) < 0 {
		return nil, errors.New("error allocating input picture")
	}

	if e.h = C.x264_encoder_open(&e.params); e.h == nil {
		return nil, errors.New("error opening encoder")
	}

	return e, nil
}

func (e *Encoder) Encode(pts int64, f frame.Frame) error {
	//TODO: make this a Source channel
	var nal *C.x264_nal_t
	var i_nal C.int
	var frame_size C.int

	if int8F, ok := f.(frame.TypedFrame[uint8]); ok {
		e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeLuma()[0]))
		e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCb()[0]))
		e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&int8F.GetNativeCr()[0]))
	} else if int16F, ok := f.(frame.TypedFrame[uint16]); ok {
		e.pictureIn.img.plane[0] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeLuma()[0]))
		e.pictureIn.img.plane[1] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCb()[0]))
		e.pictureIn.img.plane[2] = (*C.uint8_t)(unsafe.Pointer(&int16F.GetNativeCr()[0]))
	}

	//cleanup pointers
	defer func() {
		e.pictureIn.img.plane[0] = nil
		e.pictureIn.img.plane[1] = nil
		e.pictureIn.img.plane[2] = nil

	}()
	defer runtime.KeepAlive(f)

	e.pictureIn.i_pts = C.int64_t(pts)

	if frame_size = C.x264_encoder_encode(e.h, &nal, &i_nal, e.pictureIn, e.pictureOut); frame_size < 0 {
		return errors.New("error encoding frame")
	}

	if frame_size == 0 {
		return nil
	}

	buf := unsafe.Slice((*byte)(nal.p_payload), int(frame_size))
	if _, err := e.w.Write(buf); err != nil {
		return err
	}

	return nil
}

func (e *Encoder) Flush() error {
	var nal *C.x264_nal_t
	var i_nal C.int
	var frame_size C.int

	for C.x264_encoder_delayed_frames(e.h) > 0 {
		if frame_size = C.x264_encoder_encode(e.h, &nal, &i_nal, nil, e.pictureOut); frame_size <= 0 {
			return errors.New("error encoding frame")
		}

		buf := unsafe.Slice((*byte)(nal.p_payload), int(frame_size))
		if _, err := e.w.Write(buf); err != nil {
			return err
		}

	}

	return nil
}

func (e *Encoder) Close() {
	if e.cleaned.Swap(true) == false {
		if e.h != nil {
			C.x264_encoder_close(e.h)
			e.h = nil
		}
		if e.pictureIn != nil {
			C.x264_picture_clean(e.pictureIn)
			C.free(unsafe.Pointer(e.pictureIn))
			e.pictureIn = nil
		}
		if e.pictureOut != nil {
			C.free(unsafe.Pointer(e.pictureOut))
			e.pictureOut = nil
		}
	}
}

func getSettingString(m map[string]any, name string, fallback string) string {
	if v, ok := m[name]; ok {
		if val, ok := v.(string); ok {
			return val
		}
		if val, ok := v.(int); ok {
			return strconv.Itoa(val)
		}
		if val, ok := v.(int64); ok {
			return strconv.Itoa(int(val))
		}
	}
	return fallback
}

func getSettingInt(m map[string]any, name string, fallback int) int {
	return fallback
}