package rtsp

import (
    "context"
    "fmt"
    "log"
    "net"
    "sync"
    "time"

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

    // We keep this import so the CreateServer() signature matches the original
    // repo (actioncam.go calls rtsp.CreateServer(ctx, host, port, camera)).
    "github.com/jonas-koeritz/actioncam/libipcamera"
)

// This is where your RTP relay should send H.264 packets.
// The original project uses port 5220 for the preview stream.
const defaultRTPInPort = 5220

// Server wraps the gortsplib server + the in-memory stream.
type Server struct {
    gs      *gortsplib.Server
    stream  *gortsplib.ServerStream
    media   *description.Media
    ctx     context.Context
    cancel  context.CancelFunc
    wg      sync.WaitGroup
    rtpConn *net.UDPConn
}

// handler implements the gortsplib ServerHandler* interfaces.
type handler struct {
    srv *Server
}

// --- RTSP callbacks (multi-client capable) ---

// OnDescribe: clients ask what the stream looks like (SDP).
func (h *handler) OnDescribe(ctx *gortsplib.ServerHandlerOnDescribeCtx) (*base.Response, *gortsplib.ServerStream, error) {
    log.Printf("RTSP DESCRIBE from %v path=%s", ctx.Conn.NetConn().RemoteAddr(), ctx.Path)

    return &base.Response{
        StatusCode: base.StatusOK,
    }, h.srv.stream, nil
}

// OnSetup: client wants to SUBSCRIBE to the existing stream.
func (h *handler) OnSetup(ctx *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, *gortsplib.ServerStream, error) {
    log.Printf("RTSP SETUP from %v path=%s", ctx.Conn.NetConn().RemoteAddr(), ctx.Path)

    return &base.Response{
        StatusCode: base.StatusOK,
    }, h.srv.stream, nil
}

// OnPlay: client starts receiving packets.
func (h *handler) OnPlay(ctx *gortsplib.ServerHandlerOnPlayCtx) (*base.Response, error) {
    log.Printf("RTSP PLAY from %v path=%s", ctx.Conn.NetConn().RemoteAddr(), ctx.Path)

    return &base.Response{
        StatusCode: base.StatusOK,
    }, nil
}

// (Optional) For logging / debugging:
func (h *handler) OnConnOpen(ctx *gortsplib.ServerHandlerOnConnOpenCtx) {
    log.Printf("RTSP connection opened: %v", ctx.Conn.NetConn().RemoteAddr())
}

func (h *handler) OnConnClose(ctx *gortsplib.ServerHandlerOnConnCloseCtx) {
    log.Printf("RTSP connection closed: %v (err=%v)", ctx.Conn.NetConn().RemoteAddr(), ctx.Error)
}

// --- Public API ---

// CreateServer creates and starts a RTSP server that serves *one* H.264 video
// stream, backed by a single UDP source (127.0.0.1:defaultRTPInPort).
//
// It keeps the original function signature from the repo:
//   rtsp.CreateServer(applicationContext, host, port, camera)
// but the camera is NOT used directly here – you’re expected to start the
// RTP relay separately so that it forwards H.264 RTP packets into
// 127.0.0.1:defaultRTPInPort.
//
// Multiple RTSP clients are automatically supported: all of them read from
// the same gortsplib.ServerStream.
func CreateServer(parentCtx context.Context, host string, port int, _ *libipcamera.Camera) (*Server, error) {
    ctx, cancel := context.WithCancel(parentCtx)

    // Build an in-memory H.264 stream description suitable for gortsplib.
    stream, media, err := newH264Stream()
    if err != nil {
        cancel()
        return nil, fmt.Errorf("create H264 stream: %w", err)
    }

    // Listen for incoming RTP packets from the RTP relay.
    rtpAddr := &net.UDPAddr{
        IP:   net.ParseIP("127.0.0.1"),
        Port: defaultRTPInPort,
    }

    rtpConn, err := net.ListenUDP("udp", rtpAddr)
    if err != nil {
        cancel()
        return nil, fmt.Errorf("listen udp %v: %w", rtpAddr, err)
    }

    srv := &Server{
        stream:  stream,
        media:   media,
        ctx:     ctx,
        cancel:  cancel,
        rtpConn: rtpConn,
    }

    h := &handler{srv: srv}

    // Configure gortsplib server.
    gs := &gortsplib.Server{
        Handler:     h,
        RTSPAddress: fmt.Sprintf("%s:%d", host, port),
        // If you also want UDP transport for clients, set UDPRTPAddress/UDPRTCPAddress here.
        // For most use cases, TCP (interleaved over RTSP) is fine.
    }
    srv.gs = gs

    // Start the UDP → RTSP pump.
    srv.wg.Add(1)
    go srv.rtpPump()

    // Start RTSP server in the background.
    go func() {
        log.Printf("RTSP server listening on rtsp://%s:%d/ (gortsplib)", host, port)

        if err := gs.StartAndWait(); err != nil {
            log.Printf("RTSP server stopped with error: %v", err)
        }
        cancel()
    }()

    return srv, nil
}

// Close shuts down the RTSP server and the RTP pump.
func (s *Server) Close() {
    s.cancel()
    if s.gs != nil {
        s.gs.Close()
    }
    if s.rtpConn != nil {
        _ = s.rtpConn.Close()
    }
    s.wg.Wait()
    if s.stream != nil {
        s.stream.Close()
    }
}

// --- internals ---

// newH264Stream builds a single-video-track ServerStream with a valid clock rate.
//
// NOTE:
//  * SPS / PPS here are "generic valid" values, not tuned to your camera.
//  * For best results, you can parse the real SPS/PPS from your camera.sdp and
//    drop them in here.
func newH264Stream() (*gortsplib.ServerStream, *description.Media, error) {
    // Generic baseline-profile SPS/PPS. They just need to be *valid* so that
    // the library knows the clock rate and doesn't panic ("non-positive interval
    // for NewTicker") :contentReference[oaicite:1]{index=1}
    h264 := &format.H264{
        PayloadTyp: 96,
        SPS: []byte{
            0x67, 0x42, 0xC0, 0x1F, 0x96, 0x54, 0x05, 0x01, 0xED, 0x00, 0xF0, 0x88, 0x45, 0x80,
        },
        PPS: []byte{
            0x68, 0xCE, 0x38, 0x80,
        },
        PacketizationMode: 1,
    }

    media := &description.Media{
        Type: description.MediaTypeVideo,
        // You could also set media.Control if you want a specific track URL.
        Formats: []format.Format{h264},
    }

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

    stream := &gortsplib.ServerStream{
        Desc: desc,
    }
    if err := stream.Initialize(); err != nil {
        return nil, nil, err
    }

    return stream, media, nil
}

// rtpPump reads RTP packets from UDP and pushes them into the gortsplib stream.
// Every connected RTSP client gets the same packets (multi-client fan-out).
func (s *Server) rtpPump() {
    defer s.wg.Done()

    buf := make([]byte, 2048)

    for {
        select {
        case <-s.ctx.Done():
            return
        default:
        }

        // Avoid blocking forever so we can react to shutdown.
        _ = s.rtpConn.SetReadDeadline(time.Now().Add(2 * time.Second))

        n, _, err := s.rtpConn.ReadFromUDP(buf)
        if err != nil {
            if ne, ok := err.(net.Error); ok && ne.Timeout() {
                continue
            }
            // If the context is done, exit quietly.
            if s.ctx.Err() != nil {
                return
            }

            log.Printf("RTP read error: %v", err)
            continue
        }

        var pkt rtp.Packet
        if err := pkt.Unmarshal(buf[:n]); err != nil {
            // Ignore malformed packets.
            continue
        }

        if err := s.stream.WritePacketRTP(s.media, &pkt); err != nil {
            // This is non-fatal; clients can disconnect at any time.
            log.Printf("WritePacketRTP error: %v", err)
        }
    }
}