var upgrader = websocket.Upgrader{} // use default options

func ws(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer conn.Close()
	for {
		mt, message, err := conn.ReadMessage()
		log.Printf("recv: %s", message)
		conn.WriteMessage(mt, message)
	}
}


func main() {
	http.HandleFunc("/", ws) // Will create a goroutine to execute the function
        log.Fatal(http.ListenAndServe(":8080", nil))
}

// All methods can be mapped to Linux epoll commands (http://man7.org/linux/man-pages/man7/epoll.7.html)
 
// Create an epoll with initial space for fds.
// Since Linux 2.6.8, the size argument is ignored, but must be greater than zero for backward compatible.
func EpollCreate(size int) (epfd int, err error) {......}

// Create an epoll with indicating flag. 
// If flag is 0, then, other than the fact that the obsolete size argument is dropped
func EpollCreate1(flag int) (epfd int, err error) {......}

// EpollCtl is used to add, modify, or remove entries in the interest list of the epoll.
// Valid values for the op argument are: EPOLL_CTL_ADD, EPOLL_CTL_MOD, and EPOLL_CTL_DEL.
// EpollEvent is an object telling the epoll what events we want to monitor for the fd.
// E.g. unix.EpollEvent{Events: unix.POLLIN | unix.POLLHUP, Fd: int32(fd)}
func EpollCtl(epfd int, op int, fd int, event *EpollEvent) (err error) {......}

// EpollWait call system call waits for events on the epoll.
// Ready events will be injected to events
func EpollWait(epfd int, events []EpollEvent, msec int) (n int, err error) {......}

type epoll struct {
	fd          int
	connections map[int]net.Conn // net.Conn implements io.ReadWriter
	lock        *sync.RWMutex
}

func MkEpoll() (*epoll, error) {
	fd, err := unix.EpollCreate1(0)
	if err != nil {
		return nil, err
	}
	return &epoll{
		fd:          fd,
		lock:        &sync.RWMutex{},
		connections: make(map[int]net.Conn),
	}, nil
}

func socketFD(conn net.Conn) int {
	tcpConn := reflect.Indirect(reflect.ValueOf(conn)).FieldByName("conn")
	fdVal := tcpConn.FieldByName("fd")
	pfdVal := reflect.Indirect(fdVal).FieldByName("pfd")
	return int(pfdVal.FieldByName("Sysfd").Int())
}

func (e *epoll) Add(conn net.Conn) error {
	// Extract file descriptor associated with the connection
	fd := socketFD(conn)
        event := unix.EpollEvent{Events: unix.POLLIN | unix.POLLHUP, Fd: int32(fd)}
	err := unix.EpollCtl(e.fd, syscall.EPOLL_CTL_ADD, fd, &event)
	if err != nil {
		return err
	}
	e.connections[fd] = conn // connections are protected by RWMutex
	return nil
}

func (e *epoll) Remove(conn net.Conn) error {
	fd := socketFD(conn)
	err := unix.EpollCtl(e.fd, syscall.EPOLL_CTL_DEL, fd, nil)
	if err != nil {
		return err
	}
	delete(e.connections, fd)
	return nil
}

func (e *epoll) Wait() ([]net.Conn, error) {
	events := make([]unix.EpollEvent, 100) // The max size of events retrieving from epoll
	n, err := unix.EpollWait(e.fd, events, 100)
	if err != nil {
		return nil, err
	}
	var connections []net.Conn
	for i := 0; i < n; i++ {
		conn := e.connections[int(events[i].Fd)]
		connections = append(connections, conn)
	}
	return connections, nil
}

var epoller *epoll

func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, _, _, err := ws.UpgradeHTTP(r, w) // Upgrade connection
	if err != nil {
		return
	}
        // When conn is added to epoller, this goroutine is done.
	if err := epoller.Add(conn); err != nil {
		log.Printf("Failed to add connection %v", err)
		conn.Close()
	}
}

func main() {
	// Increase resources limitations......
        // Enable pprof hooks.......

	// Start epoll
	var err error
	epoller, err = MkEpoll()
	if err != nil {
		panic(err)
	}

	go HandleConnection()

	http.HandleFunc("/", wsHandler)
	if err := http.ListenAndServe("0.0.0.0:8000", nil); err != nil {
		log.Fatal(err)
	}
}

func HandleConnection() {
	for {
		connections, err := epoller.Wait()
		if err != nil {
			log.Printf("Failed to epoll wait %v", err)
			continue
		}
		for _, conn := range connections {
			if conn == nil {
				break
			}
                        // ReadClientData reads next data message from conn
			if msg, _, err := wsutil.ReadClientData(conn); err != nil {
				if err := epoller.Remove(conn); err != nil {
					log.Printf("Failed to remove %v", err)
				}
				conn.Close()
			} else {
				log.Printf("msg: %s", string(msg))
			}
		}
	}