package rtsp

import (
	"context"
	"fmt"
	"log"
	"sync"

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

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

type Server struct {
	stream *gortsplib.ServerStream
	media  *description.Media
	gs     *gortsplib.Server
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

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

    // Create the RTSP server instance first
    h := &handler{}
    gs := &gortsplib.Server{
        Handler:     h,
        UDPRTPAddress:  ":8000",
        UDPRTCPAddress: ":8001",
        RTSPAddress: fmt.Sprintf("%s:%d", host, port),
    }

    // Now create the H264 stream bound to this server
    stream, media, err := newH264Stream(gs)
    if err != nil {
        cancel()
        return nil, fmt.Errorf("create H264 stream: %w", err)
    }

    // Final server struct assembly
    srv := &Server{
        ctx:    ctx,
        cancel: cancel,
        gs:     gs,
        stream: stream,
        media:  media,
    }
    h.srv = srv // link handler to server instance

    // Start pumping RTP frames from relay into RTSP
    srv.wg.Add(1)
    go srv.relayPump(relay)

    // Start the RTSP server
    go func() {
        log.Printf("RTSP server ready on rtsp://%s:%d/", host, port)
        if err := gs.StartAndWait(); err != nil {
            log.Printf("RTSP stopped: %v", err)
        }
        cancel()
    }()

    return srv, nil
}

func (s *Server) Close() {
	s.cancel()
	if s.gs != nil {
		s.gs.Close()
	}
	s.wg.Wait()
	if s.stream != nil {
		s.stream.Close()
	}
}

// Relay pump: consumes RTP from channel
func (s *Server) relayPump(relay *libipcamera.RTPRelay) {
	defer s.wg.Done()
	for {
		select {
		case <-s.ctx.Done():
			return
		case frame, ok := <-relay.Output:
			if !ok {
				return
			}
			// Build RTP packet
			var pkt rtp.Packet
			if err := pkt.Unmarshal(frame.Payload); err != nil {
				continue
			}
			s.stream.WritePacketRTP(s.media, &pkt)
		}
	}
}

// Handler below…
type handler struct{ srv *Server }

func (h *handler) OnDescribe(ctx *gortsplib.ServerHandlerOnDescribeCtx) (*base.Response, *gortsplib.ServerStream, error) {
	return &base.Response{StatusCode: base.StatusOK}, h.srv.stream, nil
}

func (h *handler) OnSetup(ctx *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, *gortsplib.ServerStream, error) {
	return &base.Response{StatusCode: base.StatusOK}, h.srv.stream, nil
}

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

// Build simple baseline H264 SDP
func newH264Stream(gs *gortsplib.Server) (*gortsplib.ServerStream, *description.Media, error) {
    h264 := &format.H264{
        PayloadTyp:        96,
        PacketizationMode: 1,
        SPS:               []byte{0x67, 0x42, 0xC0, 0x1F, 0x96, 0x54, 0x05, 0x01, 0xED, 0x00, 0xF0, 0x88, 0x45, 0x80},
        PPS:               []byte{0x68, 0xCE, 0x38, 0x80},
    }

    media := &description.Media{
        Type:    description.MediaTypeVideo,
        Formats: []format.Format{h264},
    }

    desc := &description.Session{
        Medias: []*description.Media{media},
    }

    // REQUIRED: link stream to the running server
    stream := gortsplib.NewServerStream(gs, desc)

    return stream, media, nil
}