package main

import (
	"context"
	"log"
	"net/http"
	"time"

	"nhooyr.io/websocket"
	"nhooyr.io/websocket/wsjson"
)

type Server struct {
	OnOpen    func(conn *websocket.Conn)
	OnClose   func(conn *websocket.Conn, err error)
	OnMessage func(conn *websocket.Conn, msg map[string]any)
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
		InsecureSkipVerify: true,
	})
	if err != nil {
		log.Println("accept error:", err)
		return
	}
	defer conn.CloseNow()

	if s.OnOpen != nil {
		s.OnOpen(conn)
	}

	ctx := r.Context()

	var readErr error
	for {
		var msg map[string]any
		if readErr = wsjson.Read(ctx, conn, &msg); readErr != nil {
			break
		}
		if s.OnMessage != nil {
			s.OnMessage(conn, msg)
		}
	}

	if s.OnClose != nil {
		s.OnClose(conn, readErr)
	}

	conn.Close(websocket.StatusNormalClosure, "done")
}

func main() {
	srv := &Server{
		OnOpen: func(conn *websocket.Conn) {
			log.Println("client connected")
		},
		OnClose: func(conn *websocket.Conn, err error) {
			log.Println("client disconnected:", err)
		},
		OnMessage: func(conn *websocket.Conn, msg map[string]any) {
			log.Printf("received: %v\n", msg)
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			if err := wsjson.Write(ctx, conn, map[string]any{"echo": msg}); err != nil {
				log.Println("write error:", err)
			}
		},
	}

	http.Handle("/ws", srv)
	log.Println("server listening on :8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}