package main

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

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"
)

type wsServer struct {
	OnOpen    func(c *Client)
	OnClose   func(c *Client, err error)
	OnMessage func(c *Client, msg map[string]any)
}

var (
	clients []*Client
	mu      sync.Mutex
)

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

	client := &Client{conn: conn}
	mu.Lock()
	clients = append(clients, client)
	mu.Unlock()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

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

	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(client, msg)
		}
	}

	cancel() // cancel before OnClose so any in-flight queries are canceled first

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

	mu.Lock()
	for i, c := range clients {
		if c == client {
			clients[i] = clients[len(clients)-1]
			clients = clients[:len(clients)-1]
			break
		}
	}
	mu.Unlock()

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

func sendAndCloseIfFails(conn *websocket.Conn, message map[string]any) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := wsjson.Write(ctx, conn, message); err != nil {
		conn.Close(websocket.StatusGoingAway, "Write error")
	}
}

func sendToAllExceptAndCloseIfFails(client *Client, message map[string]any) {
	for _, other := range clients {
		if other != client {
			sendAndCloseIfFails(other.conn, message)
		}
	}
}

func handleUnauthenticatedMessage(client *Client, msg map[string]any) {
	token := msg["token"].(string)
	user, err := GetUserFromToken(token)
	if err != nil {
		client.conn.Close(websocket.StatusPolicyViolation, "invalid token")
		return
	}
	client.User = &user
	sendAndCloseIfFails(client.conn, map[string]any{
		"authAs": user.Name,
	})
	log.Println("New User authenticated as: " + user.Name)
}

func handleAuthenticatedMessage(client *Client, msg map[string]any) {
	message := msg["message"].(string)
	if message == "" {
		sendAndCloseIfFails(client.conn, map[string]any{
			"error": "no message",
		})
		return
	}

	sendToAllExceptAndCloseIfFails(client, map[string]any{
		"username": client.User.Name,
		"message":  message,
	})
}