package source

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	mxl "github.com/qvest-digital/go-mxl/mxl"
)

type flowDef struct {
	FrameWidth  int    `json:"frame_width"`
	FrameHeight int    `json:"frame_height"`
	MediaType   string `json:"media_type"`
	Colorspace  string `json:"colorspace"`
	GrainRate   struct {
		Numerator   int64 `json:"numerator"`
		Denominator int64 `json:"denominator"`
	} `json:"grain_rate"`
}

type Frame struct {
	Index   uint64
	Width   uint32
	Height  uint32
	Stride  uint32
	Size    uint32
	Invalid bool
	Payload []byte
}

type Source struct {
	inst   *mxl.Instance
	reader *mxl.Reader
	info   mxl.FlowInfo
	def    string
	rate   mxl.Rational
	stride uint32
	width  uint32
	height uint32
	idx    uint64
}

func Open(domain, flowID string) (*Source, error) {
	inst, err := mxl.NewInstance(domain, "")
	if err != nil {
		return nil, fmt.Errorf("NewInstance: %w", err)
	}
	r, err := inst.NewReader(flowID)
	if err != nil {
		inst.Close()
		return nil, fmt.Errorf("NewReader: %w", err)
	}
	info, err := r.Info()
	if err != nil {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("Info: %w", err)
	}
	def, err := inst.FlowDef(flowID)
	if err != nil {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("FlowDef: %w", err)
	}
	var fd flowDef
	if err := json.Unmarshal([]byte(def), &fd); err != nil {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("parse flow def: %w", err)
	}
	if fd.FrameWidth == 0 || fd.FrameHeight == 0 {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("flow has no video dimensions (not a video flow?)")
	}
	rate := info.Config.Common.GrainRate
	idx := mxl.CurrentIndex(rate)
	if idx == mxl.UndefinedIndex {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("invalid grain rate: %d/%d", rate.Num, rate.Den)
	}
	return &Source{
		inst:   inst,
		reader: r,
		info:   info,
		def:    def,
		rate:   info.Config.Common.GrainRate,
		stride: info.Config.Discrete.SliceSizes[0],
		width:  uint32(fd.FrameWidth),
		height: uint32(fd.FrameHeight),
		idx:    idx,
	}, nil
}

func (s *Source) Close() error {
	_ = s.reader.Close()
	return s.inst.Close()
}

func (s *Source) Next(timeout time.Duration) (Frame, error) {
	return s.NextCtx(context.Background(), timeout)
}

func (s *Source) NextCtx(ctx context.Context, timeout time.Duration) (Frame, error) {
	for {
		select {
		case <-ctx.Done():
			return Frame{}, ctx.Err()
		default:
		}
		g, err := s.reader.GetGrain(s.idx, timeout)
		switch {
		case err == nil:
			f := Frame{
				Index:   g.Index,
				Width:   s.width,
				Height:  s.height,
				Stride:  s.stride,
				Size:    g.GrainSize,
				Invalid: g.Invalid(),
				Payload: g.Payload,
			}
			s.idx++
			return f, nil
		case errors.Is(err, mxl.ErrTimeout):
			s.idx = mxl.CurrentIndex(s.rate)
		case errors.Is(err, mxl.ErrOutOfRangeEarly):
			time.Sleep(10 * time.Millisecond)
		case errors.Is(err, mxl.ErrOutOfRangeLate):
			s.idx = mxl.CurrentIndex(s.rate)
		default:
			return Frame{}, fmt.Errorf("GetGrain: %w", err)
		}
	}
}

func (s *Source) FlowDef() string        { return s.def }
func (s *Source) Rate() mxl.Rational     { return s.rate }
func (s *Source) Stride() uint32         { return s.stride }
func (s *Source) Width() uint32          { return s.width }
func (s *Source) Height() uint32         { return s.height }
func (s *Source) GrainCount() uint32     { return s.info.Config.Discrete.GrainCount }
func (s *Source) Format() mxl.DataFormat { return s.info.Config.Common.Format }

type AudioSource struct {
	inst  *mxl.Instance
	r     *mxl.Reader
	info  mxl.FlowInfo
	rate  mxl.Rational
	chans uint64
	idx   uint64
}

type AudioFrame struct {
	Index       uint64
	SampleCount uint64
	Channels    uint64
	Samples     [][]byte // per-channel byte slices (F32, deinterleaved)
}

func OpenAudio(domain, flowID string) (*AudioSource, error) {
	inst, err := mxl.NewInstance(domain, "")
	if err != nil {
		return nil, fmt.Errorf("NewInstance: %w", err)
	}
	r, err := inst.NewReader(flowID)
	if err != nil {
		inst.Close()
		return nil, fmt.Errorf("NewReader: %w", err)
	}
	info, err := r.Info()
	if err != nil {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("Info: %w", err)
	}
	if info.Config.Common.Format.IsDiscrete() {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("flow is discrete (not audio)")
	}
	idx := info.Runtime.HeadIndex
	if idx == 0 {
		r.Close()
		inst.Close()
		return nil, fmt.Errorf("flow has no head yet (no producer?)")
	}
	return &AudioSource{
		inst:  inst,
		r:     r,
		info:  info,
		rate:  info.Config.Common.GrainRate,
		chans: uint64(info.Config.Continuous.ChannelCount),
		idx:   idx,
	}, nil
}

func (s *AudioSource) NextAudio(ctx context.Context, batch uint64, timeout time.Duration) (AudioFrame, error) {
	for {
		select {
		case <-ctx.Done():
			return AudioFrame{}, ctx.Err()
		default:
		}
		v, err := s.r.GetSamples(s.idx, int(batch), timeout)
		switch {
		case err == nil:
			samples := make([][]byte, s.chans)
			for ch := uint64(0); ch < s.chans; ch++ {
				f1, f2, _ := v.ChannelFragments(ch)
				if len(f2) > 0 {
					samples[ch] = append(f1, f2...)
				} else {
					samples[ch] = f1
				}
			}
			f := AudioFrame{
				Index:       s.idx,
				SampleCount: batch,
				Channels:    s.chans,
				Samples:     samples,
			}
			s.idx += batch
			return f, nil
		case errors.Is(err, mxl.ErrOutOfRangeEarly):
			select {
			case <-time.After(10 * time.Millisecond):
			case <-ctx.Done():
				return AudioFrame{}, ctx.Err()
			}
		case errors.Is(err, mxl.ErrOutOfRangeLate):
			rt, rerr := s.r.Runtime()
			if rerr != nil {
				return AudioFrame{}, fmt.Errorf("Runtime: %w", rerr)
			}
			s.idx = rt.HeadIndex
		default:
			return AudioFrame{}, fmt.Errorf("GetSamples: %w", err)
		}
	}
}

func (s *AudioSource) Close() error {
	_ = s.r.Close()
	return s.inst.Close()
}

func (s *AudioSource) Rate() mxl.Rational { return s.rate }
func (s *AudioSource) Channels() uint64   { return s.chans }

type SyncSource struct {
	inst                  *mxl.Instance
	vr                    *mxl.Reader
	ar                    *mxl.Reader
	group                 *mxl.SyncGroup
	rate                  mxl.Rational // video rate
	aRate                 mxl.Rational
	chans                 uint64
	idx                   uint64
	width, height, stride uint32
}

func OpenSync(domain, videoFlow, audioFlow string) (*SyncSource, error) {
	inst, err := mxl.NewInstance(domain, "")
	if err != nil {
		return nil, fmt.Errorf("NewInstance: %w", err)
	}

	vr, err := inst.NewReader(videoFlow)
	if err != nil {
		inst.Close()
		return nil, fmt.Errorf("NewReader(video): %w", err)
	}
	ar, err := inst.NewReader(audioFlow)
	if err != nil {
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("NewReader(audio): %w", err)
	}

	vInfo, err := vr.Info()
	if err != nil {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("Info(video): %w", err)
	}
	if !vInfo.Config.Common.Format.IsDiscrete() {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("video flow is not discrete")
	}
	aInfo, err := ar.Info()
	if err != nil {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("Info(audio): %w", err)
	}
	if aInfo.Config.Common.Format.IsDiscrete() {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("audio flow is not continuous")
	}

	def, err := inst.FlowDef(videoFlow)
	if err != nil {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("FlowDef: %w", err)
	}
	var fd flowDef
	if err := json.Unmarshal([]byte(def), &fd); err != nil {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("parse flow def: %w", err)
	}

	vRate := vInfo.Config.Common.GrainRate
	idx := mxl.CurrentIndex(vRate)
	if idx == mxl.UndefinedIndex {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("invalid video rate %d/%d", vRate.Num, vRate.Den)
	}

	group, err := inst.NewSyncGroup()
	if err != nil {
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("NewSyncGroup: %w", err)
	}
	if err := group.AddReader(vr); err != nil {
		group.Close()
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("AddReader(video): %w", err)
	}
	if err := group.AddReader(ar); err != nil {
		group.Close()
		ar.Close()
		vr.Close()
		inst.Close()
		return nil, fmt.Errorf("AddReader(audio): %w", err)
	}

	return &SyncSource{
		inst:   inst,
		vr:     vr,
		ar:     ar,
		group:  group,
		rate:   vRate,
		aRate:  aInfo.Config.Common.GrainRate,
		chans:  uint64(aInfo.Config.Continuous.ChannelCount),
		idx:    idx,
		width:  uint32(fd.FrameWidth),
		height: uint32(fd.FrameHeight),
		stride: vInfo.Config.Discrete.SliceSizes[0],
	}, nil
}

func (s *SyncSource) Close() error {
	_ = s.group.Close()
	_ = s.ar.Close()
	_ = s.vr.Close()
	return s.inst.Close()
}

// NextSync reads both at a synced timestamp. Returns video Frame + audio AudioFrame
func (s *SyncSource) NextSync(ctx context.Context, audioBatch uint64, timeout time.Duration) (Frame, AudioFrame, error) {
	var timeouts int
	for {
		select {
		case <-ctx.Done():
			return Frame{}, AudioFrame{}, ctx.Err()
		default:
		}

		ts := mxl.IndexToTimestamp(s.rate, s.idx)
		err := s.group.WaitForDataAt(ts, timeout)
		switch {
		case err == nil:
			// read video
			g, gerr := s.vr.GetGrain(s.idx, 50*time.Millisecond)
			if gerr != nil {
				if errors.Is(gerr, mxl.ErrFlowInvalid) {
					s.idx = mxl.CurrentIndex(s.rate)
					continue
				}
				if errors.Is(gerr, mxl.ErrOutOfRangeLate) {
					s.idx = mxl.CurrentIndex(s.rate)
					continue
				}
				if errors.Is(gerr, mxl.ErrOutOfRangeEarly) {
					select {
					case <-time.After(5 * time.Millisecond):
					case <-ctx.Done():
						return Frame{}, AudioFrame{}, ctx.Err()
					}
					continue
				}
				return Frame{}, AudioFrame{}, fmt.Errorf("GetGrain: %w", gerr)
			}
			// read audio at the same timestamp
			aIdx := mxl.TimestampToIndex(s.aRate, ts)
			av, aerr := s.ar.GetSamples(aIdx, int(audioBatch), 50*time.Millisecond)
			vFrame := Frame{
				Index: g.Index, Width: s.width, Height: s.height,
				Stride: s.stride, Size: g.GrainSize,
				Invalid: g.Invalid(), Payload: g.Payload,
			}
			s.idx++

			var aFrame AudioFrame
			if aerr == nil {
				samples := make([][]byte, s.chans)
				for ch := uint64(0); ch < s.chans; ch++ {
					f1, f2, _ := av.ChannelFragments(ch)
					if len(f2) > 0 {
						samples[ch] = append(f1, f2...)
					} else {
						samples[ch] = f1
					}
				}
				aFrame = AudioFrame{
					Index: aIdx, SampleCount: audioBatch,
					Channels: s.chans, Samples: samples,
				}
			}
			// even if audio failed, video returns
			return vFrame, aFrame, nil
		case errors.Is(err, mxl.ErrTimeout), errors.Is(err, mxl.ErrOutOfRangeEarly), errors.Is(err, mxl.ErrOutOfRangeLate):
			timeouts++
			if timeouts > 10 {
				timeouts = 0
				s.idx = mxl.CurrentIndex(s.rate)
				return Frame{}, AudioFrame{}, fmt.Errorf("sync: feeds not responding")
			}
			s.idx = mxl.CurrentIndex(s.rate)
			select {
			case <-time.After(5 * time.Millisecond):
			case <-ctx.Done():
				return Frame{}, AudioFrame{}, ctx.Err()
			}
		default:
			return Frame{}, AudioFrame{}, fmt.Errorf("WaitForDataAt: %w", err)
		}
	}
}

func (s *SyncSource) Width() uint32           { return s.width }
func (s *SyncSource) Height() uint32          { return s.height }
func (s *SyncSource) Stride() uint32          { return s.stride }
func (s *SyncSource) Rate() mxl.Rational      { return s.rate }
func (s *SyncSource) AudioRate() mxl.Rational { return s.aRate }
func (s *SyncSource) Channels() uint64        { return s.chans }