package main

import (
	"context"
	_ "embed"
	"errors"
	"flag"
	"fmt"
	"net/http/httptest"
	"os"
	"strings"
	"time"

	omni "git.yuntuo.hk/k2safe/go-omni"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/ws"
	"git.yuntuo.hk/k2safe/go-omni/ws/hub"
	wsomni "git.yuntuo.hk/k2safe/go-omni/ws/omni"
)

//go:embed index.html
var page []byte

func main() {
	serve := flag.Bool("serve", false, "serve the browser demo on 127.0.0.1:8088")
	flag.Parse()
	if err := run(*serve); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run(serve bool) (runErr error) {
	application := omni.New(omni.WithLogger(false))
	defer func() { runErr = errors.Join(runErr, application.Close()) }()
	connections, err := hub.New(hub.Config{})
	if err != nil {
		return err
	}
	application.OnShutdown(connections.Close)
	server, err := ws.NewServer(ws.ServerConfig{})
	if err != nil {
		return err
	}
	if err = wsomni.Mount(application, "/ws", server, ws.Handlers{
		OnOpen: func(_ context.Context, s *ws.Session) error {
			if err := connections.Add(s); err != nil {
				return err
			}
			return connections.Join(s.ID(), "demo")
		},
		OnMessage: func(ctx context.Context, s *ws.Session, m ws.Message) error {
			if m.Type == ws.Text && string(m.Data) == "broadcast" {
				result, err := connections.SendGroup("demo", ws.Message{Type: ws.Text, Data: []byte("hello group")})
				if err != nil {
					return err
				}
				if len(result.Errors) > 0 {
					return s.SendJSON(ctx, map[string]int{"queued": result.Queued, "failed": len(result.Errors)})
				}
				return nil
			}
			return s.Send(ctx, m)
		},
	}); err != nil {
		return err
	}
	application.GET("/", func(c *httpx.Context) error {
		c.Raw().Header().Set("Content-Type", "text/html; charset=utf-8")
		_, err := c.Raw().Write(page)
		return err
	})
	if serve {
		fmt.Println("Open http://127.0.0.1:8088 in two browser tabs; Ctrl+C stops cleanly.")
		return application.Run("127.0.0.1:8088")
	}
	httpServer := httptest.NewServer(application)
	defer httpServer.Close()
	client, err := ws.NewClient(ws.ClientConfig{ProxyMode: "direct"})
	if err != nil {
		return err
	}
	defer func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		runErr = errors.Join(runErr, client.Shutdown(ctx))
	}()
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	replies := make(chan string, 2)
	session, err := client.Dial(ctx, "ws"+strings.TrimPrefix(httpServer.URL, "http")+"/ws", ws.Handlers{OnMessage: func(_ context.Context, _ *ws.Session, m ws.Message) error { replies <- string(m.Data); return nil }})
	if err != nil {
		return err
	}
	for _, message := range []string{"hello websocket", "broadcast"} {
		if err = session.Send(ctx, ws.Message{Type: ws.Text, Data: []byte(message)}); err != nil {
			return err
		}
		select {
		case reply := <-replies:
			fmt.Println(reply)
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	return nil
}