package udp

import (
	"context"
	"fmt"
	"net"
	"runtime"
	"sync"
)

type Server struct {
	bindAddr     string
	routingTable *RoutingTable

	conns  []*net.UDPConn
	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc
}

func NewServer(bindAddr string) *Server {
	ctx, cancel := context.WithCancel(context.Background())

	return &Server{
		bindAddr:     bindAddr,
		routingTable: NewRoutingTable(),
		ctx:          ctx,
		cancel:       cancel,
	}
}

func (s *Server) Router() *RoutingTable {
	return s.routingTable
}

func (s *Server) Run() error {
	if s.bindAddr == "" {
		return fmt.Errorf("bind address is empty")
	}

	workerCount := runtime.NumCPU()

	conns, err := listenUDP(s.bindAddr, workerCount)
	if err != nil {
		return fmt.Errorf("cannot listen on address: %w", err)
	}
	if len(conns) == 0 {
		return fmt.Errorf("no UDP connections created")
	}
	s.conns = conns

	for _, conn := range s.conns {
		_ = conn.SetReadBuffer(8 * 1024 * 1024)
		_ = conn.SetWriteBuffer(8 * 1024 * 1024)
	}

	fmt.Println("[udp] listening on", s.bindAddr, "with", len(s.conns), "worker(s)")

	for workerID, conn := range s.conns {
		s.wg.Add(1)
		go s.workerLoop(workerID, conn)
	}

	return nil
}

func (s *Server) Stop() {
	s.cancel()
	for _, c := range s.conns {
		_ = c.Close()
	}
	s.wg.Wait()
}

func (s *Server) workerLoop(workerID int, conn *net.UDPConn) {
	defer s.wg.Done()

	buf := make([]byte, 1500)
	fmt.Println("[udp] worker", workerID, "started")

	for {
		select {
		case <-s.ctx.Done():
			fmt.Println("[udp] worker", workerID, "stopped")
			return
		default:
			n, addr, err := conn.ReadFromUDP(buf)
			if err != nil {
				if s.ctx.Err() != nil {
					return
				}
				if ne, ok := err.(net.Error); ok && ne.Timeout() {
					continue
				}
				fmt.Printf("[udp] worker %d: read error: %v\n", workerID, err)
				continue
			}

			// Le worker qui lit est aussi celui qui traite et qui écrit.
			s.handlePacket(conn, buf[:n], addr)
		}
	}
}

// handlePacket reçoit maintenant directement la conn du worker.
// Aucun hop supplémentaire : le paquet reste dans la goroutine/conn du worker.
func (s *Server) handlePacket(conn *net.UDPConn, data []byte, addr *net.UDPAddr) {
	if len(data) == 0 {
		return
	}

	pt := PacketType(data[0])

	switch pt {
	case PacketTypePing:
		// exemple simple : echo du ping
		_, _ = conn.WriteToUDP([]byte{byte(PacketTypePing)}, addr)

	case PacketTypeConnect:
		if len(data) < 2 {
			return
		}
		channelID := string(data[1:])
		s.routingTable.Add(channelID, addr)

	case PacketTypeDisconnect:
		if len(data) < 2 {
			return
		}
		channelID := string(data[1:])
		s.routingTable.Remove(channelID, addr)

	case PacketTypeVoiceData:
		if len(data) < 2 {
			return
		}
		channelID := string(data[1:]) // à adapter selon ton vrai format
		recipients := s.routingTable.GetAddrs(channelID)

		for _, dst := range recipients {
			// optionnel: ne pas renvoyer à la source
			if dst.IP.Equal(addr.IP) && dst.Port == addr.Port {
				continue
			}
			_, _ = conn.WriteToUDP(data, dst)
		}

	default:
		// type inconnu -> ignore
	}
}