package main

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/pion/webrtc/v4"
	"github.com/x0gp/server/internal/control"
	"github.com/x0gp/server/internal/lobby"
	"github.com/x0gp/server/internal/transport"
)

type driverSession struct {
	driverID    string
	deviceID    uint8
	pc          *webrtc.PeerConnection
	videoDC     *webrtc.DataChannel
	dataDC      *webrtc.DataChannel
	lastCommand string
	cancelCtx   context.CancelFunc
}

type WebRTCService struct {
	logger     *slog.Logger
	engine     *control.Engine
	lobbySvc   *lobby.Service
	videoSvc   *UDPVideoService
	sessions   map[string]*driverSession
	sessionsMu sync.RWMutex
}

func NewWebRTCService(logger *slog.Logger, e *control.Engine, lobbySvc *lobby.Service, videoSvc *UDPVideoService) *WebRTCService {
	return &WebRTCService{
		logger:   logger,
		engine:   e,
		lobbySvc: lobbySvc,
		videoSvc: videoSvc,
		sessions: make(map[string]*driverSession),
	}
}

func (s *WebRTCService) Start(ctx context.Context, wg *sync.WaitGroup) {
	wg.Add(1)
	go func() {
		defer wg.Done()
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				s.CloseAll()
				return
			case <-ticker.C:
				s.broadcastTelemetry()
			}
		}
	}()
}

// ConnectHandler godoc
// @Summary      Establish WebRTC connection for driver-car 1-to-1 proxying
// @Description  Accepts SDP Offer, negotiates peer connection, and returns SDP Answer. Routes video and control data via DataChannels.
// @Tags         webrtc
// @Accept       json
// @Produce      json
// @Param        body body transport.WebRTCConnectRequest true "SDP Offer details"
// @Success      200 {object} transport.WebRTCConnectResponse "SDP Answer negotiation success"
// @Router       /api/webrtc/connect [post]
func (s *WebRTCService) ConnectHandler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "Only POST is allowed")
			return
		}

		var req transport.WebRTCConnectRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, "bad_request", "Invalid request JSON")
			return
		}

		if req.DriverID == "" {
			writeError(w, http.StatusBadRequest, "bad_request", "driver_id is required")
			return
		}

		// Resolve device ID for this driver
		deviceID := uint8(1) // default fallback
		if driver, err := s.lobbySvc.GetDriver(req.DriverID); err == nil && driver.DeviceID != nil {
			deviceID = uint8(*driver.DeviceID)
		} else {
			s.logger.Warn("webrtc client connected but no physical device_id assigned in lobbySvc. Defaulting to device 1", "driver_id", req.DriverID)
		}

		// Create PeerConnection
		config := webrtc.Configuration{
			ICEServers: []webrtc.ICEServer{
				{
					URLs: []string{"stun:stun.l.google.com:19302"},
				},
			},
		}

		pc, err := webrtc.NewPeerConnection(config)
		if err != nil {
			s.logger.Error("failed to create webrtc peer connection", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Prepare context for background video pumping
		connCtx, connCancel := context.WithCancel(context.Background())

		session := &driverSession{
			driverID:  req.DriverID,
			deviceID:  deviceID,
			pc:        pc,
			cancelCtx: connCancel,
		}

		// Handle Data Channels
		pc.OnDataChannel(func(d *webrtc.DataChannel) {
			s.logger.Info("received remote data channel", "label", d.Label(), "driver_id", req.DriverID)
			s.sessionsMu.Lock()
			if d.Label() == "video" {
				session.videoDC = d
			} else if d.Label() == "data" {
				session.dataDC = d
			}
			s.sessionsMu.Unlock()

			d.OnOpen(func() {
				s.logger.Info("data channel opened", "label", d.Label(), "driver_id", req.DriverID)
				if d.Label() == "video" {
					// Start pumping video frames from UDPVideoService to the data channel
					go s.pumpVideo(connCtx, session)
				}
			})

			d.OnMessage(func(msg webrtc.DataChannelMessage) {
				if d.Label() == "data" {
					s.handleDataChannelMessage(session, msg.Data)
				}
			})
		})

		pc.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
			s.logger.Info("webrtc connection state changed", "state", state.String(), "driver_id", req.DriverID)
			if state == webrtc.PeerConnectionStateClosed || state == webrtc.PeerConnectionStateFailed {
				s.CloseSession(req.DriverID)
			}
		})

		// Set the remote SessionDescription
		err = pc.SetRemoteDescription(webrtc.SessionDescription{
			Type: webrtc.SDPTypeOffer,
			SDP:  req.SDP,
		})
		if err != nil {
			pc.Close()
			connCancel()
			s.logger.Error("failed to set remote description", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Create Answer
		answer, err := pc.CreateAnswer(nil)
		if err != nil {
			pc.Close()
			connCancel()
			s.logger.Error("failed to create answer", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Set Local Description
		err = pc.SetLocalDescription(answer)
		if err != nil {
			pc.Close()
			connCancel()
			s.logger.Error("failed to set local description", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Wait for ICE Gathering to complete (Vanilla ICE)
		gatherComplete := webrtc.GatheringCompletePromise(pc)
		<-gatherComplete

		// Save Session
		s.sessionsMu.Lock()
		// Clean up existing session if any
		if existing, ok := s.sessions[req.DriverID]; ok {
			existing.pc.Close()
			existing.cancelCtx()
		}
		s.sessions[req.DriverID] = session
		s.sessionsMu.Unlock()

		localDesc := pc.LocalDescription()
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(transport.WebRTCConnectResponse{
			SDP:  localDesc.SDP,
			Type: "answer",
		})
	}
}

func (s *WebRTCService) pumpVideo(ctx context.Context, session *driverSession) {
	s.logger.Info("starting WebRTC video pump for driver", "driver_id", session.driverID, "device_id", session.deviceID)

	frameCh := make(chan []byte, 10)
	s.videoSvc.registerClient(session.deviceID, frameCh)
	defer func() {
		s.videoSvc.unregisterClient(session.deviceID, frameCh)
		s.logger.Info("stopping WebRTC video pump for driver", "driver_id", session.driverID)
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case frame, ok := <-frameCh:
			if !ok {
				return
			}
			s.sessionsMu.RLock()
			dc := session.videoDC
			s.sessionsMu.RUnlock()

			if dc != nil && dc.ReadyState() == webrtc.DataChannelStateOpen {
				if err := dc.Send(frame); err != nil {
					s.logger.Debug("failed to send video frame over datachannel", "driver_id", session.driverID, "err", err)
				}
			}
		}
	}
}

func (s *WebRTCService) handleDataChannelMessage(session *driverSession, data []byte) {
	var payload map[string]any
	if err := json.Unmarshal(data, &payload); err != nil {
		s.logger.Warn("received invalid JSON on data channel", "driver_id", session.driverID, "err", err)
		return
	}

	// 1. Steering & Throttle inputs
	in := transport.InputState{}
	hasInput := false
	if v, ok := payload["throttle"].(float64); ok {
		in.Throttle = v
		hasInput = true
	}
	if v, ok := payload["steering"].(float64); ok {
		in.Steering = v
		hasInput = true
	}
	if v, ok := payload["brake"].(float64); ok {
		in.Brake = v
		hasInput = true
	}
	if v, ok := payload["gear"].(float64); ok {
		in.Gear = int(v)
		hasInput = true
	}

	if hasInput {
		s.logger.Debug("received inputs on webrtc data channel", "driver_id", session.driverID, "throttle", in.Throttle, "steering", in.Steering)
		car := s.engine.GetCarByDriver(session.driverID)
		if car != nil {
			car.ApplyInput(in, 1.0/60.0)
		}

		// Send commands to physical car
		var cmd string
		if in.Throttle > 0.1 {
			cmd = "start"
		} else {
			cmd = "stop"
		}
		if in.Steering < -0.2 {
			cmd = "left"
		} else if in.Steering > 0.2 {
			cmd = "right"
		}

		// Resolve destination device
		var destDeviceID *uint8
		if car != nil && car.DeviceID != nil {
			val := uint8(*car.DeviceID)
			destDeviceID = &val
		} else if session.deviceID != 0 {
			// Fallback to connection-time device
			destDeviceID = &session.deviceID
		}

		if destDeviceID != nil && cmd != "" && cmd != session.lastCommand {
			s.logger.Info("webrtc control state transition", "driver_id", session.driverID, "old_cmd", session.lastCommand, "new_cmd", cmd)
			session.lastCommand = cmd
			err := s.videoSvc.SendCommand(*destDeviceID, cmd)
			if err != nil {
				s.logger.Error("failed to forward command to ESP32", "device_id", *destDeviceID, "cmd", cmd, "err", err)
			} else {
				s.logger.Info("forwarded command to ESP32 via UDP", "device_id", *destDeviceID, "cmd", cmd)
			}
		}
	}

	// Echo back message / send ACK
	s.sessionsMu.RLock()
	dc := session.dataDC
	s.sessionsMu.RUnlock()
	if dc != nil && dc.ReadyState() == webrtc.DataChannelStateOpen {
		ack, _ := json.Marshal(map[string]any{
			"type":      "ack",
			"applied":   true,
			"timestamp": time.Now().UnixMilli(),
		})
		_ = dc.SendText(string(ack))
	}
}

func (s *WebRTCService) broadcastTelemetry() {
	s.sessionsMu.RLock()
	defer s.sessionsMu.RUnlock()

	for _, session := range s.sessions {
		dc := session.dataDC
		if dc != nil && dc.ReadyState() == webrtc.DataChannelStateOpen {
			// Pull some stats/telemetry
			car := s.engine.GetCarByDriver(session.driverID)
			var speed float64
			var x, y float64
			var carID string
			if car != nil {
				speed = car.Speed
				x = car.X
				y = car.Y
				carID = car.ID
			}

			telemetry := map[string]any{
				"type":      "telemetry",
				"car_id":    carID,
				"driver_id": session.driverID,
				"speed":     speed,
				"x":         x,
				"y":         y,
				"timestamp": time.Now().UnixMilli(),
			}

			payload, _ := json.Marshal(telemetry)
			_ = dc.SendText(string(payload))
		}
	}
}

func (s *WebRTCService) CloseSession(driverID string) {
	s.sessionsMu.Lock()
	defer s.sessionsMu.Unlock()

	if session, ok := s.sessions[driverID]; ok {
		s.logger.Info("closing WebRTC session for driver", "driver_id", driverID)
		session.pc.Close()
		session.cancelCtx()
		delete(s.sessions, driverID)
	}
}

func (s *WebRTCService) CloseAll() {
	s.sessionsMu.Lock()
	defer s.sessionsMu.Unlock()

	for driverID, session := range s.sessions {
		s.logger.Info("closing WebRTC session (shutdown)", "driver_id", driverID)
		session.pc.Close()
		session.cancelCtx()
	}
	s.sessions = make(map[string]*driverSession)
}