package rtsp

import (
	"context"
	"log"
	"sync"
	"strconv"

	"github.com/bluenviron/gortsplib/v5"
	"github.com/bluenviron/gortsplib/v5/pkg/base"
	"github.com/bluenviron/gortsplib/v5/pkg/description"
	"github.com/bluenviron/gortsplib/v5/pkg/format"
	"github.com/pion/rtp"

	"github.com/jonas-koeritz/actioncam/libipcamera"
)

type Handler struct {
	server *gortsplib.Server
	stream *gortsplib.ServerStream
	mu     sync.RWMutex
}

func (h *Handler) OnDescribe(_ *gortsplib.ServerHandlerOnDescribeCtx) (*base.Response, *gortsplib.ServerStream, error) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return &base.Response{StatusCode: base.StatusOK}, h.stream, nil
}

func (h *Handler) OnSetup(_ *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, *gortsplib.ServerStream, error) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return &base.Response{StatusCode: base.StatusOK}, h.stream, nil
}

func (h *Handler) OnPlay(_ *gortsplib.ServerHandlerOnPlayCtx) (*base.Response, error) {
	return &base.Response{StatusCode: base.StatusOK}, nil
}

type Server struct {
	server *gortsplib.Server
	stream *gortsplib.ServerStream
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

func CreateServer(ctx context.Context, host string, port int, relay *libipcamera.RTPRelay) *Server {
	cctx, cancel := context.WithCancel(ctx)

	h := &Handler{}
	srv := &Server{cancel: cancel}

	h.server = &gortsplib.Server{
		Handler:        h,
		RTSPAddress:    host + ":" + strconv.Itoa(port),
		UDPRTPAddress:  ":8000",
		UDPRTCPAddress: ":8001",
	}

	if err := h.server.Start(); err != nil {
		panic(err)
	}

	h.mu.Lock()
	desc := &description.Session{
		Medias: []*description.Media{{
			Type: description.MediaTypeVideo,
			Formats: []format.Format{&format.H264{
				PayloadTyp:        96,
				PacketizationMode: 1,
			}},
		}},
	}
	h.stream = &gortsplib.ServerStream{Server: h.server, Desc: desc}
	if err := h.stream.Initialize(); err != nil {
		panic(err)
	}
	srv.server = h.server
	srv.stream = h.stream
	h.mu.Unlock()

	// Pump frames -> RTP
	srv.wg.Add(1)
	go func() {
		defer srv.wg.Done()
		var seq uint16
		var ts uint32
		for {
			select {
			case <-cctx.Done():
				return
			case frame, ok := <-relay.Frames:
				if !ok {
					return
				}
				pkt := &rtp.Packet{
					Header: rtp.Header{
						Version:        2,
						Marker:         true,
						PayloadType:    96,
						SequenceNumber: seq,
						Timestamp:      ts,
					},
					Payload: frame.Data,
				}
				seq++
				ts += 3600
				h.stream.WritePacketRTP(h.stream.Desc.Medias[0], pkt)
			}
		}
	}()

	log.Printf("RTSP server ready at rtsp://%s:%d/", host, port)
	return srv
}

func (s *Server) Stop() {
	s.cancel()
	s.server.Close()
	s.stream.Close()
	s.wg.Wait()
}