// Package hls contains HLS utilities.
package hls

import (
	"encoding/hex"
	"errors"
	"fmt"
	"reflect"
	"slices"
	"strings"

	"github.com/bluenviron/gohlslib/v2"
	"github.com/bluenviron/gohlslib/v2/pkg/codecs"
	"github.com/bluenviron/gortsplib/v5/pkg/description"
	"github.com/bluenviron/gortsplib/v5/pkg/format"
	"github.com/bluenviron/mediacommon/v2/pkg/codecs/flac"
	"github.com/bluenviron/mediacommon/v2/pkg/codecs/mpeg4audio"

	"github.com/bluenviron/mediamtx/internal/formatlabel"
	"github.com/bluenviron/mediamtx/internal/logger"
	"github.com/bluenviron/mediamtx/internal/stream"
	"github.com/bluenviron/mediamtx/internal/unit"
)

// ErrNoSupportedCodecs is returned by FromStream when there are no supported codecs.
var ErrNoSupportedCodecs = errors.New(
	"the stream doesn't contain any supported codec, which are currently AV1, VP9, H265, H264, Opus, MPEG-4 Audio, KLV")

func findFormatAndIndexInMedia(media *description.Media, forma any) int {
	for i, forma2 := range media.Formats {
		if reflect.TypeOf(forma2) == reflect.TypeOf(forma).Elem() {
			reflect.ValueOf(forma).Elem().Set(reflect.ValueOf(forma2))
			return i
		}
	}

	return -1
}

func findFormatAndIndexes(d *description.Session, forma any) (*description.Media, int, int) {
	for i, media := range d.Medias {
		j := findFormatAndIndexInMedia(media, forma)
		if j >= 0 {
			return media, i, j
		}
	}
	return nil, -1, -1
}

func setupVideoTrack(
	origDesc *description.Session,
	outDesc *description.Session,
	r *stream.Reader,
	muxer *gohlslib.Muxer,
) {
	addTrack := func(
		media *description.Media,
		forma format.Format,
		track *gohlslib.Track,
		onData stream.OnDataFunc,
	) {
		muxer.Tracks = append(muxer.Tracks, track)
		r.OnData(media, forma, onData)
	}

	var videoFormatAV1 *format.AV1
	videoMedia, _, _ := findFormatAndIndexes(origDesc, &videoFormatAV1)

	if videoFormatAV1 != nil {
		track := &gohlslib.Track{
			Codec:     &codecs.AV1{},
			ClockRate: videoFormatAV1.ClockRate(),
		}

		addTrack(
			videoMedia,
			videoFormatAV1,
			track,
			func(u *unit.Unit) error {
				if u.NilPayload() {
					return nil
				}

				err := muxer.WriteAV1(
					track,
					u.NTP,
					u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
					u.Payload.(unit.PayloadAV1))
				if err != nil {
					return fmt.Errorf("muxer error: %w", err)
				}

				return nil
			})

		return
	}

	var videoFormatVP9 *format.VP9
	videoMedia, _, _ = findFormatAndIndexes(origDesc, &videoFormatVP9)

	if videoFormatVP9 != nil {
		track := &gohlslib.Track{
			Codec:     &codecs.VP9{},
			ClockRate: videoFormatVP9.ClockRate(),
		}

		addTrack(
			videoMedia,
			videoFormatVP9,
			track,
			func(u *unit.Unit) error {
				if u.NilPayload() {
					return nil
				}

				err := muxer.WriteVP9(
					track,
					u.NTP,
					u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
					u.Payload.(unit.PayloadVP9))
				if err != nil {
					return fmt.Errorf("muxer error: %w", err)
				}

				return nil
			})

		return
	}

	var videoFormatH265 *format.H265
	videoMedia, i, j := findFormatAndIndexes(origDesc, &videoFormatH265)

	if videoFormatH265 != nil {
		outFormat := outDesc.Medias[i].Formats[j].(*format.H265)

		track := &gohlslib.Track{
			Codec: &codecs.H265{
				VPS: outFormat.VPS,
				SPS: outFormat.SPS,
				PPS: outFormat.PPS,
			},
			ClockRate: videoFormatH265.ClockRate(),
		}

		addTrack(
			videoMedia,
			videoFormatH265,
			track,
			func(u *unit.Unit) error {
				if u.NilPayload() {
					return nil
				}

				err := muxer.WriteH265(
					track,
					u.NTP,
					u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
					u.Payload.(unit.PayloadH265))
				if err != nil {
					return fmt.Errorf("muxer error: %w", err)
				}

				return nil
			})

		return
	}

	var videoFormatH264 *format.H264
	videoMedia, i, j = findFormatAndIndexes(origDesc, &videoFormatH264)

	if videoFormatH264 != nil {
		outFormat := outDesc.Medias[i].Formats[j].(*format.H264)

		track := &gohlslib.Track{
			Codec: &codecs.H264{
				SPS: outFormat.SPS,
				PPS: outFormat.PPS,
			},
			ClockRate: videoFormatH264.ClockRate(),
		}

		addTrack(
			videoMedia,
			videoFormatH264,
			track,
			func(u *unit.Unit) error {
				if u.NilPayload() {
					return nil
				}

				err := muxer.WriteH264(
					track,
					u.NTP,
					u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
					u.Payload.(unit.PayloadH264))
				if err != nil {
					return fmt.Errorf("muxer error: %w", err)
				}

				return nil
			})

		return
	}
}

func setupAudioTracks(
	origDesc *description.Session,
	_ *description.Session,
	r *stream.Reader,
	muxer *gohlslib.Muxer,
) error {
	addTrack := func(
		medi *description.Media,
		forma format.Format,
		track *gohlslib.Track,
		onData stream.OnDataFunc,
	) {
		muxer.Tracks = append(muxer.Tracks, track)
		r.OnData(medi, forma, onData)
	}

	for _, media := range origDesc.Medias {
		for _, forma := range media.Formats {
			switch forma := forma.(type) {
			case *format.Opus:
				track := &gohlslib.Track{
					Codec: &codecs.Opus{
						ChannelCount: forma.ChannelCount,
					},
					ClockRate: forma.ClockRate(),
				}

				addTrack(
					media,
					forma,
					track,
					func(u *unit.Unit) error {
						err := muxer.WriteOpus(
							track,
							u.NTP,
							u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
							u.Payload.(unit.PayloadOpus))
						if err != nil {
							return fmt.Errorf("muxer error: %w", err)
						}

						return nil
					})

			case *format.Generic:
				if strings.HasPrefix(strings.ToLower(forma.RTPMap()), "flac/") {
					enc, err := hex.DecodeString(forma.FMT["streaminfo"])
					if err != nil {
						return err
					}

					var streamInfo flac.StreamInfo
					err = streamInfo.Unmarshal(enc)
					if err != nil {
						return err
					}

					track := &gohlslib.Track{
						Codec: &codecs.FLAC{
							StreamInfo: &streamInfo,
						},
						ClockRate: forma.ClockRate(),
					}

					addTrack(
						media,
						forma,
						track,
						func(u *unit.Unit) error {
							if u.NilPayload() {
								return nil
							}

							err2 := muxer.WriteFLAC(
								track,
								u.NTP,
								u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
								u.Payload.(unit.PayloadFLAC))
							if err2 != nil {
								return fmt.Errorf("muxer error: %w", err2)
							}

							return nil
						})
				}

			case *format.MPEG4Audio:
				track := &gohlslib.Track{
					Codec: &codecs.MPEG4Audio{
						Config: *forma.Config,
					},
					ClockRate: forma.ClockRate(),
				}

				addTrack(
					media,
					forma,
					track,
					func(u *unit.Unit) error {
						if u.NilPayload() {
							return nil
						}

						err := muxer.WriteMPEG4Audio(
							track,
							u.NTP,
							u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
							u.Payload.(unit.PayloadMPEG4Audio))
						if err != nil {
							return fmt.Errorf("muxer error: %w", err)
						}

						return nil
					})

			case *format.MPEG4AudioLATM:
				if !forma.CPresent {
					track := &gohlslib.Track{
						Codec: &codecs.MPEG4Audio{
							Config: *forma.StreamMuxConfig.Programs[0].Layers[0].AudioSpecificConfig,
						},
						ClockRate: forma.ClockRate(),
					}

					addTrack(
						media,
						forma,
						track,
						func(u *unit.Unit) error {
							if u.NilPayload() {
								return nil
							}

							var ame mpeg4audio.AudioMuxElement
							ame.StreamMuxConfig = forma.StreamMuxConfig
							err := ame.Unmarshal(u.Payload.(unit.PayloadMPEG4AudioLATM))
							if err != nil {
								return err
							}

							return muxer.WriteMPEG4Audio(
								track,
								u.NTP,
								u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
								[][]byte{ame.Payloads[0][0][0]})
						})
				}
			}
		}
	}

	return nil
}

func setupDataTracks(
	origDesc *description.Session,
	_ *description.Session,
	r *stream.Reader,
	muxer *gohlslib.Muxer,
) {
	addTrack := func(
		media *description.Media,
		forma format.Format,
		track *gohlslib.Track,
		onData stream.OnDataFunc,
	) {
		muxer.Tracks = append(muxer.Tracks, track)
		r.OnData(media, forma, onData)
	}

	for _, media := range origDesc.Medias {
		for _, forma := range media.Formats {
			if forma, ok := forma.(*format.KLV); ok && muxer.Variant == gohlslib.MuxerVariantMPEGTS {
				track := &gohlslib.Track{
					Codec:     &codecs.KLV{Synchronous: true},
					ClockRate: forma.ClockRate(),
				}

				addTrack(
					media,
					forma,
					track,
					func(u *unit.Unit) error {
						if u.NilPayload() {
							return nil
						}

						err := muxer.WriteKLV(
							track,
							u.NTP,
							u.PTS, // no conversion is needed since we set gohlslib.Track.ClockRate = format.ClockRate
							u.Payload.(unit.PayloadKLV))
						if err != nil {
							return fmt.Errorf("muxer error: %w", err)
						}

						return nil
					})
			}
		}
	}
}

// FromStream maps a MediaMTX stream to a HLS muxer.
func FromStream(
	origDesc *description.Session,
	outDesc *description.Session,
	r *stream.Reader,
	muxer *gohlslib.Muxer,
) error {
	setupVideoTrack(
		origDesc,
		outDesc,
		r,
		muxer,
	)

	err := setupAudioTracks(
		origDesc,
		outDesc,
		r,
		muxer,
	)
	if err != nil {
		return err
	}

	setupDataTracks(
		origDesc,
		outDesc,
		r,
		muxer,
	)

	if len(muxer.Tracks) == 0 {
		return ErrNoSupportedCodecs
	}

	setuppedFormats := r.Formats()

	n := 1
	for _, media := range origDesc.Medias {
		for _, forma := range media.Formats {
			if !slices.Contains(setuppedFormats, forma) {
				r.Parent.Log(logger.Warn, "skipping track %d (%s)", n, formatlabel.FormatToLabel(forma))
			}
			n++
		}
	}

	return nil
}