package ws

import (
	"context"
	"errors"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// ServerConfig is independent of HTTP body limits and HTTP request deadlines.
// Context owns accepted sessions; it must outlive HTTP handlers. Prepare runs
// before upgrading and may authenticate or create per-connection handlers.
type ServerConfig struct {
	Session            SessionConfig                                          `json:"session"`
	MaxConnections     int                                                    `json:"maxConnections"`
	HandshakeTimeout   time.Duration                                          `json:"handshakeTimeout"`
	AllowedOrigins     []string                                               `json:"allowedOrigins"`
	RequireOrigin      bool                                                   `json:"requireOrigin"`
	Subprotocols       []string                                               `json:"subprotocols"`
	RequireSubprotocol bool                                                   `json:"requireSubprotocol"`
	EnableCompression  bool                                                   `json:"enableCompression"`
	HandshakeRate      float64                                                `json:"handshakeRate"` // Per second; zero disables rate limiting.
	HandshakeBurst     int                                                    `json:"handshakeBurst"`
	Context            context.Context                                        `json:"-"`
	Prepare            func(context.Context, *http.Request) (Handlers, error) `json:"-"`
}

func DefaultServerConfig() ServerConfig {
	return ServerConfig{Session: DefaultSessionConfig(), MaxConnections: 4096, HandshakeTimeout: 10 * time.Second}
}

func (cfg ServerConfig) Validate() error { _, err := normalizeServerConfig(cfg); return err }

func normalizeServerConfig(cfg ServerConfig) (ServerConfig, error) {
	var err error
	cfg.Session, err = normalizeSessionConfig(cfg.Session)
	if err != nil {
		return cfg, err
	}
	if cfg.MaxConnections < 0 || cfg.MaxConnections > 1<<20 || cfg.HandshakeTimeout < 0 || cfg.HandshakeRate < 0 || cfg.HandshakeRate > 1e9 || cfg.HandshakeRate != cfg.HandshakeRate || cfg.HandshakeBurst < 0 || cfg.HandshakeBurst > 1<<20 {
		return cfg, errors.New("ws: invalid server limits")
	}
	if cfg.MaxConnections == 0 {
		cfg.MaxConnections = 4096
	}
	if cfg.HandshakeTimeout == 0 {
		cfg.HandshakeTimeout = 10 * time.Second
	}
	if cfg.HandshakeRate > 0 && cfg.HandshakeBurst == 0 {
		cfg.HandshakeBurst = 1
	}
	if cfg.HandshakeRate == 0 && cfg.HandshakeBurst != 0 {
		return cfg, errors.New("ws: handshake burst requires a rate")
	}
	if len(cfg.AllowedOrigins) > 1024 {
		return cfg, errors.New("ws: too many allowed origins")
	}
	cfg.AllowedOrigins = append([]string(nil), cfg.AllowedOrigins...)
	for i, v := range cfg.AllowedOrigins {
		o, ok := canonicalOrigin(v)
		if !ok {
			return cfg, errors.New("ws: invalid allowed origin")
		}
		cfg.AllowedOrigins[i] = o
	}
	if len(cfg.Subprotocols) > 64 {
		return cfg, errors.New("ws: too many subprotocols")
	}
	cfg.Subprotocols = append([]string(nil), cfg.Subprotocols...)
	seen := make(map[string]bool)
	for _, v := range cfg.Subprotocols {
		if !protocolToken(v) || seen[v] {
			return cfg, errors.New("ws: invalid or duplicate subprotocol")
		}
		seen[v] = true
	}
	if cfg.RequireSubprotocol && len(cfg.Subprotocols) == 0 {
		return cfg, errors.New("ws: required subprotocol list is empty")
	}
	if cfg.Context == nil {
		cfg.Context = context.Background()
	}
	return cfg, nil
}

func protocolToken(s string) bool {
	if len(s) == 0 || len(s) > 256 {
		return false
	}
	for _, r := range s {
		if r <= 32 || r >= 127 || strings.ContainsRune("()<>@,;:\\\"/[]?={}", r) {
			return false
		}
	}
	return true
}

func canonicalOrigin(raw string) (string, bool) {
	if len(raw) > 4096 || strings.ContainsAny(raw, " \t\r\n,\\?#") {
		return "", false
	}
	u, err := url.Parse(raw)
	if err != nil || u == nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || u.Opaque != "" || strings.Contains(u.Host, "*") || strings.HasSuffix(u.Host, ":") {
		return "", false
	}
	host := strings.ToLower(u.Host)
	// net/url accepts some unbracketed colon hosts. An Origin must contain a
	// serialized DNS host or a bracketed, valid IPv6 address, never that ambiguity.
	name := u.Hostname()
	if strings.Contains(name, ":") {
		if !strings.HasPrefix(u.Host, "[") || net.ParseIP(name) == nil {
			return "", false
		}
	} else {
		if strings.ContainsAny(u.Host, "[]") {
			return "", false
		}
		for _, r := range name {
			if !(r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '.' || r == '-' || r == '_') {
				return "", false
			}
		}
	}
	if port := u.Port(); port != "" {
		n, err := strconv.Atoi(port)
		if err != nil || n < 1 || n > 65535 {
			return "", false
		}
	}
	if u.Scheme == "http" && u.Port() == "80" {
		host = strings.TrimSuffix(host, ":80")
	}
	if u.Scheme == "https" && u.Port() == "443" {
		host = strings.TrimSuffix(host, ":443")
	}
	return u.Scheme + "://" + host, true
}

// RejectError lets Prepare select a public HTTP 4xx status. Error details are
// deliberately never written into the handshake response.
type RejectError struct{ StatusCode int }

func (e *RejectError) Error() string { return "ws: handshake rejected" }

type ServerStats struct {
	ActiveConnections, Handshakes int
	Accepted, Rejected            uint64
}

// Server owns all upgraded sessions, including their callbacks. Use Stop and
// Drain with lifecycle.Runtime, or Shutdown before releasing shared resources.
type Server struct {
	cfg                ServerConfig
	mu                 sync.Mutex
	stopped            bool
	active             int // includes reserved handshakes; protected by the same admission lock
	handshakes         int
	sessions           map[string]*Session
	done               chan struct{}
	doneClosed         bool
	accepted, rejected uint64
	tokens             float64
	refill             time.Time
	stopParent         func() bool
}

func NewServer(cfg ServerConfig) (*Server, error) {
	cfg, err := normalizeServerConfig(cfg)
	if err != nil {
		return nil, err
	}
	s := &Server{cfg: cfg, sessions: make(map[string]*Session), done: make(chan struct{}), tokens: float64(cfg.HandshakeBurst), refill: time.Now()}
	s.stopParent = context.AfterFunc(cfg.Context, func() { _ = s.Stop(context.Background()) })
	return s, nil
}

func (s *Server) Handler(h Handlers) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = s.Accept(w, r, h) })
}
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { _, _ = s.Accept(w, r, Handlers{}) }

func (s *Server) originAllowed(r *http.Request) bool {
	values := r.Header.Values("Origin")
	if len(values) == 0 {
		return !s.cfg.RequireOrigin
	}
	if len(values) != 1 {
		return false
	}
	origin, ok := canonicalOrigin(values[0])
	if !ok {
		return false
	}
	if len(s.cfg.AllowedOrigins) > 0 {
		for _, v := range s.cfg.AllowedOrigins {
			if v == origin {
				return true
			}
		}
		return false
	}
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	self, ok := canonicalOrigin(scheme + "://" + r.Host)
	return ok && origin == self
}

func (s *Server) admit() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.stopped || s.cfg.Context.Err() != nil {
		s.rejected++
		return http.StatusServiceUnavailable
	}
	if s.active >= s.cfg.MaxConnections {
		s.rejected++
		return http.StatusServiceUnavailable
	}
	if s.cfg.HandshakeRate > 0 {
		now := time.Now()
		s.tokens = min(float64(s.cfg.HandshakeBurst), s.tokens+now.Sub(s.refill).Seconds()*s.cfg.HandshakeRate)
		s.refill = now
		if s.tokens < 1 {
			s.rejected++
			return http.StatusTooManyRequests
		}
		s.tokens--
	}
	s.active++
	s.handshakes++
	return 0
}

func (s *Server) finishLocked() {
	if s.stopped && s.active == 0 && !s.doneClosed {
		s.doneClosed = true
		close(s.done)
	}
}

// Accept writes the handshake response (including failures). It returns after
// transferring the connection to Server; callers must not reuse the writer or
// retain request-scoped state in callbacks. Session uses ServerConfig.Context.
func (s *Server) Accept(w http.ResponseWriter, r *http.Request, h Handlers) (session *Session, retErr error) {
	if status := s.admit(); status != 0 {
		http.Error(w, http.StatusText(status), status)
		return nil, &HandshakeError{StatusCode: status}
	}
	transferred := false
	defer func() {
		if !transferred {
			s.mu.Lock()
			s.active--
			s.handshakes--
			s.rejected++
			s.finishLocked()
			s.mu.Unlock()
		}
	}()
	fail := func(status int, cause error) (*Session, error) {
		http.Error(w, http.StatusText(status), status)
		return nil, &HandshakeError{StatusCode: status, Cause: cause}
	}
	if r.Method != http.MethodGet || !websocket.IsWebSocketUpgrade(r) {
		return fail(http.StatusBadRequest, nil)
	}
	if !s.originAllowed(r) {
		return fail(http.StatusForbidden, nil)
	}
	if _, ok := w.(http.Hijacker); !ok {
		return fail(http.StatusNotImplemented, ErrUnsupportedTransport)
	}
	if s.cfg.RequireSubprotocol {
		matched := false
		for _, offered := range websocket.Subprotocols(r) {
			for _, supported := range s.cfg.Subprotocols {
				if offered == supported {
					matched = true
				}
			}
		}
		if !matched {
			return fail(http.StatusBadRequest, nil)
		}
	}
	ctx, cancel := context.WithTimeout(r.Context(), s.cfg.HandshakeTimeout)
	defer cancel()
	if s.cfg.Prepare != nil {
		var err error
		h, err = prepareSafely(s.cfg.Prepare, ctx, r)
		if err != nil {
			status := http.StatusUnauthorized
			var reject *RejectError
			if errors.As(err, &reject) && reject != nil && reject.StatusCode >= 400 && reject.StatusCode < 500 {
				status = reject.StatusCode
			}
			return fail(status, nil)
		}
	}
	if ctx.Err() != nil {
		return fail(http.StatusRequestTimeout, ctx.Err())
	}
	deadline, _ := ctx.Deadline()
	remaining := time.Until(deadline)
	if remaining <= 0 {
		return fail(http.StatusRequestTimeout, context.DeadlineExceeded)
	}
	upgrader := websocket.Upgrader{HandshakeTimeout: remaining, Subprotocols: s.cfg.Subprotocols, EnableCompression: s.cfg.EnableCompression, CheckOrigin: func(*http.Request) bool { return true }, Error: func(w http.ResponseWriter, _ *http.Request, status int, _ error) {
		http.Error(w, http.StatusText(status), status)
	}}
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, &HandshakeError{Cause: errors.New("upgrade failed")}
	}
	session = newSession(conn, s.cfg.Session, h, s.cfg.Context)
	s.mu.Lock()
	s.handshakes--
	s.sessions[session.ID()] = session
	s.accepted++
	stopped := s.stopped
	transferred = true
	s.mu.Unlock()
	session.start()
	if stopped {
		_ = session.Close(websocket.CloseGoingAway, "")
	}
	go func() {
		<-session.Done()
		s.mu.Lock()
		delete(s.sessions, session.ID())
		s.active--
		s.finishLocked()
		s.mu.Unlock()
	}()
	return session, nil
}

func prepareSafely(fn func(context.Context, *http.Request) (Handlers, error), ctx context.Context, r *http.Request) (h Handlers, err error) {
	defer func() {
		if recover() != nil {
			err = errors.New("ws: preparation failed")
		}
	}()
	return fn(ctx, r)
}

// Stop closes admission and begins bounded graceful closing of existing
// sessions. It never waits for callbacks and is safe to call from a callback.
func (s *Server) Stop(context.Context) error {
	s.mu.Lock()
	s.stopped = true
	all := make([]*Session, 0, len(s.sessions))
	for _, v := range s.sessions {
		all = append(all, v)
	}
	s.finishLocked()
	s.mu.Unlock()
	for _, v := range all {
		_ = v.Close(websocket.CloseGoingAway, "")
	}
	return nil
}

// Drain returns nil only when all handshakes, I/O and callbacks have exited.
// On timeout it aborts sockets, but does not claim blocked user callbacks ended.
func (s *Server) Drain(ctx context.Context) error {
	_ = s.Stop(ctx)
	select {
	case <-s.done:
		s.stopParent()
		return nil
	default:
	}
	select {
	case <-s.done:
		s.stopParent()
		return nil
	case <-ctx.Done():
		s.mu.Lock()
		all := make([]*Session, 0, len(s.sessions))
		for _, v := range s.sessions {
			all = append(all, v)
		}
		s.mu.Unlock()
		for _, v := range all {
			v.Abort()
		}
		return ctx.Err()
	}
}
func (s *Server) Shutdown(ctx context.Context) error { return s.Drain(ctx) }
func (s *Server) Stats() ServerStats {
	s.mu.Lock()
	defer s.mu.Unlock()
	return ServerStats{ActiveConnections: len(s.sessions), Handshakes: s.handshakes, Accepted: s.accepted, Rejected: s.rejected}
}