//go:build cgo && !disable_codec_libx264

package libx264

/*
#cgo pkg-config: x264
#include "libx264.h"
*/
import "C"
import (
	"errors"
	"fmt"
	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
	"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
}

var x264Version = fmt.Sprintf("x264 core:%d%s", C.X264_BUILD, C.GoString(C.Version()))

func Version() string {
	return x264Version
}

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

	preset := C.CString(getSettingString(settings, "preset", "medium"))
	defer C.free(unsafe.Pointer(preset))
	var tune *C.char
	if strVal := getSettingString(settings, "tune", ""); strVal != "" {
		tune = C.CString(strVal)
		defer C.free(unsafe.Pointer(tune))
	}

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

	defaultProfile := "high"

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

	}

	if properties.ColorSpace.BitDepth > 8 {
		e.params.i_csp |= C.X264_CSP_HIGH_DEPTH
		if defaultProfile == "high" {
			defaultProfile = "high10"
		}
	}

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

	e.params.i_bitdepth = C.int(properties.ColorSpace.BitDepth)

	e.params.i_width = C.int(properties.Width)
	e.params.i_height = C.int(properties.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(properties.FrameRate.Numerator)
	e.params.i_fps_den = C.uint32_t(properties.FrameRate.Denominator)
	if properties.FullColorRange {
		e.params.vui.b_fullrange = 1
	} else {
		e.params.vui.b_fullrange = 0
	}

	for k, v := range settings {
		if k == "profile" || k == "preset" || k == "tune" {
			continue
		}
		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.FormatInt(int64(val), 10))
			} else if val, ok := v.(int64); ok {
				strVal = C.CString(strconv.FormatInt(val, 10))
			} else if val, ok := v.(uint); ok {
				strVal = C.CString(strconv.FormatUint(uint64(val), 10))
			} else if val, ok := v.(uint64); ok {
				strVal = C.CString(strconv.FormatUint(val, 10))
			}

			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
		}
	}

	if e.params.rc.b_stat_read == 0 && e.params.rc.b_stat_write == 1 && C.GoString(preset) != "placebo" {
		//doing first pass
		C.x264_param_apply_fastfirstpass(&e.params)
	}

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

	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 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) EncodeStream(stream *frame.Stream) error {
	for f := range stream.Channel() {
		if err := e.Encode(f); err != nil {
			return err
		}
	}
	return e.Flush()
}

func (e *Encoder) Encode(f frame.Frame) error {
	var nal *C.x264_nal_t
	var iNal C.int
	var frameSize C.int

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

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

	e.pictureIn.i_pts = C.int64_t(f.PTS())

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

	if frameSize == 0 {
		return nil
	}

	if iNal != 1 {
		//return errors.New("more than one NAL present")
	}

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

	return nil
}

func (e *Encoder) Flush() error {
	var nal *C.x264_nal_t
	var iNal C.int
	var frameSize C.int

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

		if frameSize == 0 {
			return nil
		}

		if iNal != 1 {
			//return errors.New("more than one NAL present")
		}

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

	}

	return nil
}

func (e *Encoder) Version() string {
	return Version()
}

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
		}
	}
}

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
}