package wsserver

import (
	"bufio"
	"crypto/sha1"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"strings"
	"sync"
	"time"
)

// Frame opcodes.
const (
	opContinuation = 0x0
	opText         = 0x1
	opBinary       = 0x2
	opClose        = 0x8
	opPing         = 0x9
	opPong         = 0xa
)

const acceptGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

// Request is the parsed upgrade request.
type Request struct {
	Method string
	Path   string
	Query  url.Values
	Header map[string]string
}

// Handler receives a live WebSocket connection. It owns the connection and
// must return when the stream ends.
type Handler func(*Conn, *Request)

// Config starts a server.
type Config struct {
	// Listen is the TCP address, for example "127.0.0.1:3335".
	Listen string
	// Path is the required request path. Empty means "/nostrino".
	Path string
	// Handler runs one connection.
	Handler Handler
	// MaxMessage bounds one message. Zero means 32 MiB.
	MaxMessage int64
	// HandshakeTimeout bounds the upgrade. Zero means 15s.
	HandshakeTimeout time.Duration
	// Logf receives diagnostics; nil discards them.
	Logf func(string, ...any)
}

// Server is a running WebSocket listener.
type Server struct {
	cfg  Config
	ln   net.Listener
	done chan struct{}

	mu    sync.Mutex
	conns map[*Conn]struct{}
	wg    sync.WaitGroup
}

// Start listens and serves.
func Start(cfg Config) (*Server, error) {
	if cfg.Listen == "" {
		return nil, errors.New("wsserver: no listen address")
	}
	if cfg.Handler == nil {
		return nil, errors.New("wsserver: no handler")
	}
	if cfg.Path == "" {
		cfg.Path = "/nostrino"
	}
	if cfg.MaxMessage == 0 {
		cfg.MaxMessage = 32 << 20
	}
	if cfg.HandshakeTimeout == 0 {
		cfg.HandshakeTimeout = 15 * time.Second
	}
	if cfg.Logf == nil {
		cfg.Logf = func(string, ...any) {}
	}
	ln, err := net.Listen("tcp", cfg.Listen)
	if err != nil {
		return nil, fmt.Errorf("wsserver: listen: %w", err)
	}
	s := &Server{
		cfg:   cfg,
		ln:    ln,
		done:  make(chan struct{}),
		conns: make(map[*Conn]struct{}),
	}
	s.wg.Add(1)
	go s.acceptLoop()
	return s, nil
}

// Addr is the bound TCP address, including an ephemeral port.
func (s *Server) Addr() string { return s.ln.Addr().String() }

// Close stops the listener and every connection.
func (s *Server) Close() error {
	select {
	case <-s.done:
	default:
		close(s.done)
	}
	err := s.ln.Close()
	s.mu.Lock()
	conns := make([]*Conn, 0, len(s.conns))
	for c := range s.conns {
		conns = append(conns, c)
	}
	s.mu.Unlock()
	for _, c := range conns {
		c.Close()
	}
	s.wg.Wait()
	return err
}

func (s *Server) acceptLoop() {
	defer s.wg.Done()
	for {
		conn, err := s.ln.Accept()
		if err != nil {
			return
		}
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			s.serve(conn)
		}()
	}
}

func (s *Server) serve(raw net.Conn) {
	defer raw.Close()
	_ = raw.SetDeadline(time.Now().Add(s.cfg.HandshakeTimeout))
	br := bufio.NewReaderSize(raw, 16<<10)
	req, key, err := readHandshake(br)
	if err != nil {
		s.cfg.Logf("wsserver: handshake: %v", err)
		return
	}
	if req.Path != s.cfg.Path {
		fmt.Fprint(raw, "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
		return
	}
	accept := acceptKey(key)
	fmt.Fprintf(raw, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: %s\r\n\r\n", accept)
	_ = raw.SetDeadline(time.Time{})
	c := &Conn{conn: raw, br: br, max: s.cfg.MaxMessage}
	s.mu.Lock()
	s.conns[c] = struct{}{}
	s.mu.Unlock()
	s.cfg.Handler(c, req)
	c.Close()
	s.mu.Lock()
	delete(s.conns, c)
	s.mu.Unlock()
}

func readHandshake(br *bufio.Reader) (*Request, string, error) {
	line, err := br.ReadString('\n')
	if err != nil {
		return nil, "", err
	}
	line = strings.TrimRight(line, "\r\n")
	fields := strings.Fields(line)
	if len(fields) != 3 || !strings.EqualFold(fields[0], "GET") {
		return nil, "", errors.New("not a GET")
	}
	target := fields[1]
	path, query := target, ""
	if i := strings.IndexByte(target, '?'); i >= 0 {
		path, query = target[:i], target[i+1:]
	}
	q, err := url.ParseQuery(query)
	if err != nil {
		return nil, "", fmt.Errorf("bad query: %w", err)
	}
	header := make(map[string]string, 16)
	for i := 0; i < 64; i++ {
		line, err := br.ReadString('\n')
		if err != nil {
			return nil, "", err
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}
		name, value, ok := strings.Cut(line, ":")
		if !ok {
			return nil, "", errors.New("malformed header")
		}
		header[strings.ToLower(strings.TrimSpace(name))] = strings.TrimSpace(value)
	}
	if !headerContains(header["upgrade"], "websocket") || !headerContains(header["connection"], "upgrade") {
		return nil, "", errors.New("not an upgrade")
	}
	if header["sec-websocket-version"] != "13" {
		return nil, "", errors.New("unsupported websocket version")
	}
	key := header["sec-websocket-key"]
	if key == "" {
		return nil, "", errors.New("missing websocket key")
	}
	return &Request{Method: "GET", Path: path, Query: q, Header: header}, key, nil
}

func headerContains(value, token string) bool {
	for _, part := range strings.Split(value, ",") {
		if strings.EqualFold(strings.TrimSpace(part), token) {
			return true
		}
	}
	return false
}

func acceptKey(key string) string {
	h := sha1.Sum([]byte(key + acceptGUID))
	return base64.StdEncoding.EncodeToString(h[:])
}

// Conn is one server-side WebSocket and implements net.Conn with message
// payloads as the byte stream.
type Conn struct {
	conn net.Conn
	br   *bufio.Reader
	max  int64

	wmu  sync.Mutex
	dead bool

	buf       []byte
	closeOnce sync.Once
}

// Read returns bytes from the next data message.
func (c *Conn) Read(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	for len(c.buf) == 0 {
		msg, err := c.readMessage()
		if err != nil {
			return 0, err
		}
		c.buf = msg
	}
	n := copy(p, c.buf)
	c.buf = c.buf[n:]
	return n, nil
}

// Write sends p as one binary message.
func (c *Conn) Write(p []byte) (int, error) {
	if err := c.writeFrame(opBinary, p); err != nil {
		return 0, err
	}
	return len(p), nil
}

func (c *Conn) readMessage() ([]byte, error) {
	var message []byte
	started := false
	for {
		fin, op, payload, err := c.readFrame()
		if err != nil {
			return nil, err
		}
		switch op {
		case opPing:
			if err := c.writeFrame(opPong, payload); err != nil {
				return nil, err
			}
			continue
		case opPong:
			continue
		case opClose:
			_ = c.writeFrame(opClose, payload)
			c.closeConn()
			return nil, io.EOF
		case opText, opBinary:
			if started {
				return nil, errors.New("wsserver: data frame inside a message")
			}
			started = true
			message = append(message[:0], payload...)
		case opContinuation:
			if !started {
				return nil, errors.New("wsserver: continuation without a message")
			}
			message = append(message, payload...)
		default:
			return nil, fmt.Errorf("wsserver: unknown opcode %#x", op)
		}
		if int64(len(message)) > c.max {
			return nil, errors.New("wsserver: message too large")
		}
		if fin {
			if !started {
				return nil, errors.New("wsserver: empty message")
			}
			return message, nil
		}
	}
}

func (c *Conn) readFrame() (fin bool, op byte, payload []byte, err error) {
	var hdr [2]byte
	if _, err = io.ReadFull(c.br, hdr[:]); err != nil {
		return false, 0, nil, err
	}
	fin = hdr[0]&0x80 != 0
	op = hdr[0] & 0x0f
	masked := hdr[1]&0x80 != 0
	length := uint64(hdr[1] & 0x7f)
	switch length {
	case 126:
		var ext [2]byte
		if _, err = io.ReadFull(c.br, ext[:]); err != nil {
			return false, 0, nil, err
		}
		length = uint64(binary.BigEndian.Uint16(ext[:]))
	case 127:
		var ext [8]byte
		if _, err = io.ReadFull(c.br, ext[:]); err != nil {
			return false, 0, nil, err
		}
		length = binary.BigEndian.Uint64(ext[:])
		if length>>63 != 0 {
			return false, 0, nil, errors.New("wsserver: negative length")
		}
	}
	if length > uint64(c.max) {
		return false, 0, nil, errors.New("wsserver: message too large")
	}
	if op >= opClose && (!fin || length > 125) {
		return false, 0, nil, errors.New("wsserver: malformed control frame")
	}
	var mask [4]byte
	if masked {
		if _, err = io.ReadFull(c.br, mask[:]); err != nil {
			return false, 0, nil, err
		}
	}
	payload = make([]byte, length)
	if _, err = io.ReadFull(c.br, payload); err != nil {
		return false, 0, nil, err
	}
	if masked {
		for i := range payload {
			payload[i] ^= mask[i&3]
		}
	}
	return fin, op, payload, nil
}

func (c *Conn) writeFrame(op byte, payload []byte) error {
	c.wmu.Lock()
	defer c.wmu.Unlock()
	if c.dead {
		return io.ErrClosedPipe
	}
	var hdr [10]byte
	hdr[0] = 0x80 | op
	n := len(payload)
	var hlen int
	switch {
	case n < 126:
		hdr[1] = byte(n)
		hlen = 2
	case n <= 0xffff:
		hdr[1] = 126
		binary.BigEndian.PutUint16(hdr[2:4], uint16(n))
		hlen = 4
	default:
		hdr[1] = 127
		binary.BigEndian.PutUint64(hdr[2:10], uint64(n))
		hlen = 10
	}
	buf := make([]byte, hlen+n)
	copy(buf, hdr[:hlen])
	copy(buf[hlen:], payload)
	_, err := c.conn.Write(buf)
	return err
}

// Close sends a close frame and closes the socket.
func (c *Conn) Close() error {
	var err error
	c.closeOnce.Do(func() {
		_ = c.writeFrame(opClose, nil)
		err = c.closeConn()
	})
	return err
}

func (c *Conn) closeConn() error {
	c.wmu.Lock()
	c.dead = true
	c.wmu.Unlock()
	return c.conn.Close()
}

// LocalAddr implements net.Conn.
func (c *Conn) LocalAddr() net.Addr { return c.conn.LocalAddr() }

// RemoteAddr implements net.Conn.
func (c *Conn) RemoteAddr() net.Addr { return c.conn.RemoteAddr() }

// SetDeadline implements net.Conn.
func (c *Conn) SetDeadline(t time.Time) error { return c.conn.SetDeadline(t) }

// SetReadDeadline implements net.Conn.
func (c *Conn) SetReadDeadline(t time.Time) error { return c.conn.SetReadDeadline(t) }

// SetWriteDeadline implements net.Conn.
func (c *Conn) SetWriteDeadline(t time.Time) error { return c.conn.SetWriteDeadline(t) }