package udp

import (
	"fmt"
	"net"
	"sync"
	"time"
)

type Server struct {
	conn        *net.UDPConn
	connections map[string]*Connection
	mutex       sync.RWMutex
	handler     PacketHandler
	running     bool
}

type PacketHandler interface {
	HandlePacket(conn *Connection, packet *ApplicationPacket)
}

func NewServer(addr string, handler PacketHandler) (*Server, error) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return nil, err
	}

	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return nil, err
	}

	return &Server{
		conn:        conn,
		connections: make(map[string]*Connection),
		handler:     handler,
	}, nil
}

func (s *Server) Start() error {
	s.running = true

	// Start connection timeout checker
	go s.timeoutChecker()

	// Main packet receive loop
	buffer := make([]byte, 2048)
	for s.running {
		n, addr, err := s.conn.ReadFromUDP(buffer)
		if err != nil {
			if s.running {
				fmt.Printf("UDP read error: %v\n", err)
			}
			continue
		}

		go s.handlePacket(buffer[:n], addr)
	}

	return nil
}

func (s *Server) Stop() {
	s.running = false
	s.conn.Close()
}

func (s *Server) handlePacket(data []byte, addr *net.UDPAddr) {
	if len(data) < 2 {
		return
	}

	connKey := addr.String()

	s.mutex.Lock()
	conn, exists := s.connections[connKey]
	if !exists {
		conn = NewConnection(addr, s.conn, s.handler)
		s.connections[connKey] = conn
	}
	s.mutex.Unlock()

	conn.ProcessPacket(data)
}

func (s *Server) timeoutChecker() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		if !s.running {
			return
		}

		now := time.Now()
		s.mutex.Lock()
		for key, conn := range s.connections {
			if now.Sub(conn.lastPacketTime) > 45*time.Second {
				conn.Close()
				delete(s.connections, key)
			}
		}
		s.mutex.Unlock()
	}
}