//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/color"
	"git.gammaspectra.live/S.O.N.G/Ignite/frame"
	"git.gammaspectra.live/S.O.N.G/Ignite/utilities"
	"io"
	"runtime"
	"runtime/cgo"
	"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
	resourcePinner runtime.Pinner
	logger         utilities.Logger
}

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

func Version() string {
	return x264Version
}

func Build() int {
	return int(C.X264_BUILD)
}

const (
	LogLevelNone    int = C.X264_LOG_NONE
	LogLevelError   int = C.X264_LOG_ERROR
	LogLevelWarning int = C.X264_LOG_WARNING
	LogLevelInfo    int = C.X264_LOG_INFO
	LogLevelDebug   int = C.X264_LOG_DEBUG
)

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

	preset := C.CString(utilities.GetSettingString(settings, "preset", "medium"))
	defer C.free(unsafe.Pointer(preset))
	var tune *C.char
	if strVal := utilities.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(utilities.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)

	if properties.VFR {
		e.params.b_vfr_input = 1
	} else {
		e.params.b_vfr_input = 0
	}
	e.params.b_repeat_headers = 1
	e.params.b_annexb = 1

	timeBase := properties.TimeBase()

	e.params.i_timebase_num = C.uint32_t(timeBase.Numerator)
	e.params.i_timebase_den = C.uint32_t(timeBase.Denominator)
	if properties.FullColorRange {
		e.params.vui.b_fullrange = 1
	} else {
		e.params.vui.b_fullrange = 0
	}

	switch properties.ColorSpace.ChromaSamplePosition {
	case color.ChromaSamplePositionUnspecified:
		e.params.vui.i_chroma_loc = 0 //?
	case color.ChromaSamplePositionLeft:
		e.params.vui.i_chroma_loc = 0
	case color.ChromaSamplePositionCenter:
		e.params.vui.i_chroma_loc = 1
	case color.ChromaSamplePositionTopLeft:
		e.params.vui.i_chroma_loc = 2
	case color.ChromaSamplePositionTop:
		e.params.vui.i_chroma_loc = 3
	case color.ChromaSamplePositionBottomLeft:
		e.params.vui.i_chroma_loc = 4
	case color.ChromaSamplePositionBottom:
		e.params.vui.i_chroma_loc = 5
	}

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

	e.pictureIn = &C.x264_picture_t{}
	e.resourcePinner.Pin(e.pictureIn)

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

	e.params.p_log_private = unsafe.Pointer(cgo.NewHandle(e))
	C.SetLogCallback(&e.params)

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

	if err := e.encoderHeaders(); err != nil {
		return e, err
	}

	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 {
	luma := f.GetLuma()
	cb := f.GetCb()
	cr := f.GetCr()
	copy(unsafe.Slice((*byte)(e.pictureIn.img.plane[0]), len(luma)), luma)
	copy(unsafe.Slice((*byte)(e.pictureIn.img.plane[1]), len(cb)), cb)
	copy(unsafe.Slice((*byte)(e.pictureIn.img.plane[2]), len(cr)), cr)

	e.pictureIn.i_pts = C.int64_t(f.PTS())
	//TODO: read flags for statistic tracking

	return e.encoderEncode(e.pictureIn)
}

func (e *Encoder) encoderHeaders() error {
	if e.params.b_repeat_headers > 0 {
		return nil
	}

	var nal *C.x264_nal_t
	var iNal C.int
	var frameSize C.int

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

	if frameSize == 0 || iNal == 0 {
		return nil
	}

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

	// All NAL payloads are sequential to each other in memory
	buf := unsafe.Slice((*byte)(nal.p_payload), int(frameSize))
	if _, err := e.w.Write(buf); err != nil {
		return err
	}

	return nil
}

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

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

	if frameSize == 0 || iNal == 0 {
		return nil
	}

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

	// All NAL payloads are sequential to each other in memory
	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 {
	for C.x264_encoder_delayed_frames(e.h) > 0 {
		if err := e.encoderEncode(nil); err != nil {
			return err
		}
	}

	return nil
}

func (e *Encoder) Log(logLevel int, message string) {
	var prefix string
	switch logLevel {
	case C.X264_LOG_ERROR:
		prefix = "error"
	case C.X264_LOG_WARNING:
		prefix = "warning"
	case C.X264_LOG_INFO:
		prefix = "info"
	case C.X264_LOG_DEBUG:
		prefix = "debug"
	default:
		prefix = "unknown"
	}
	e.logger.Printf("x264 [%s]: %s", prefix, message)
}

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)
			e.pictureIn = nil
		}
		if e.params.p_log_private != nil {
			(cgo.Handle)(e.params.p_log_private).Delete()
			e.params.p_log_private = nil
		}
		e.resourcePinner.Unpin()
	}
}

//export logCallback
func logCallback(logLevel C.int, message *C.char, size C.size_t, data unsafe.Pointer) {
	(cgo.Handle)(data).Value().(*Encoder).Log(int(logLevel), C.GoStringN(message, C.int(size)))
}