package main

import (
	"context"
	"crypto/tls"
	"log"
	"net/http"
	"strings"
	"sync"

	"github.com/caddyserver/certmagic"
	"github.com/quic-go/quic-go/http3"
	"github.com/quic-go/webtransport-go"
)

type Room struct {
	conns map[*webtransport.Session]bool
	mu    sync.Mutex
}

type Server struct {
	rooms map[string]*Room
	mu    sync.Mutex
}

func NewServer() *Server {
	return &Server{
		rooms: make(map[string]*Room),
	}
}

func (s *Server) getRoom(name string) *Room {
	s.mu.Lock()
	defer s.mu.Unlock()

	if room, ok := s.rooms[name]; ok {
		return room
	}

	room := &Room{
		conns: make(map[*webtransport.Session]bool),
	}
	s.rooms[name] = room
	return room
}

func (s *Server) handleSession(roomName string, sess *webtransport.Session) {
	room := s.getRoom(roomName)

	room.mu.Lock()
	room.conns[sess] = true
	room.mu.Unlock()

	log.Println("joined room:", roomName)

	defer func() {
		room.mu.Lock()
		delete(room.conns, sess)
		room.mu.Unlock()
		log.Println("left room:", roomName)
	}()

	for {
		stream, err := sess.AcceptStream(context.Background())
		if err != nil {
			return
		}
		go s.handleStream(room, sess, stream)
	}
}

func (s *Server) handleStream(room *Room, sender *webtransport.Session, stream *webtransport.Stream) {
	buf := make([]byte, 4096)

	for {
		n, err := stream.Read(buf)
		if err != nil {
			return
		}

		room.mu.Lock()
		for conn := range room.conns {
			if conn == sender {
				continue
			}
			go func(c *webtransport.Session, data []byte) {
				out, err := c.OpenStream()
				if err != nil {
					return
				}
				defer out.Close()
				out.Write(data)
			}(conn, append([]byte(nil), buf[:n]...))
		}
		room.mu.Unlock()
	}
}

func main() {
	domain := "qgo-server.quizer.space"
	email := "serverussnap@outlook.com"

	certmagic.DefaultACME.Email = email
	certmagic.DefaultACME.Agreed = true
	certmagic.Default.Storage = &certmagic.FileStorage{
		Path: "/root/.local/share/certmagic",
	}

	tlsConf, err := certmagic.TLS([]string{domain})
	if err != nil {
		log.Fatal(err)
	}

	// Обязательно и h3, и http/1.1
	tlsConf.NextProtos = []string{"h3", "http/1.1"}
	tlsConf.MinVersion = tls.VersionTLS13

	server := NewServer()

	mux := http.NewServeMux()

	// WebTransport server поверх HTTP/3
	wtServer := &webtransport.Server{
		H3: &http3.Server{
			Addr:      ":443",
			TLSConfig: tlsConf,
			Handler:   mux,
		},
	}

	mux.HandleFunc("/room/", func(w http.ResponseWriter, r *http.Request) {
		roomName := strings.TrimPrefix(r.URL.Path, "/room/")

		sess, err := wtServer.Upgrade(w, r)
		if err != nil {
			log.Println("upgrade error:", err)
			return
		}

		go server.handleSession(roomName, sess)
	})

	log.Println("Relay running on https://" + domain + "/room/{room}")

	// Запуск HTTP/3 + WebTransport (UDP 443)
	go func() {
		if err := wtServer.ListenAndServe(); err != nil {
			log.Fatal(err)
		}
	}()

	// Отдельный HTTPS сервер (TCP 443) для ACME challenge
	httpServer := &http.Server{
		Addr:      ":443",
		TLSConfig: tlsConf,
		Handler:   mux,
	}

	log.Fatal(httpServer.ListenAndServeTLS("", ""))
}