package qq_obv11

// This file implements a minimal RFC 6455 WebSocket (text/binary frames,
// ping/pong, close, per-message client masking) using only the standard
// library. The project pins no third-party WebSocket dependency (go.mod must
// not change), so the OneBot v11 adapter carries its own tiny frame codec. It
// is deliberately narrow: it supports what OneBot v11 needs — UTF-8 JSON text
// frames carried over a single upgraded TCP connection — and nothing more.

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

// WebSocket control/data opcodes (RFC 6455 section 5.2).
const (
	opContinuation = 0x0
	opText         = 0x1
	opBinary       = 0x2
	opClose        = 0x8
	opPing         = 0x9
	opPong         = 0xA
)

// wsGUID is the fixed handshake GUID from RFC 6455 section 1.3.
const wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

// maxFramePayload caps a single WebSocket frame so a malicious peer cannot make
// the adapter buffer an unbounded payload. Chat/JSON payloads are tiny; this is
// generous.
const maxFramePayload = 10 << 20 // 10 MiB

// DialContextFunc dials a raw TCP (or TLS) connection. It is injectable so
// tests can route the handshake through a fake server; the default is a
// net.Dialer.
type DialContextFunc func(ctx context.Context, network, address string) (net.Conn, error)

// wsConn is a single WebSocket connection over an upgraded TCP stream. It is
// safe for one concurrent reader plus concurrent writers (writes are mutex
// serialized); exactly one goroutine may call ReadMessage at a time.
type wsConn struct {
	conn      net.Conn
	br        *bufio.Reader
	bw        *bufio.Writer
	isClient  bool
	writeMu   sync.Mutex
	closeOnce sync.Once
}

// newWSConn wraps an upgraded connection. br and bw are the buffered reader and
// writer from the handshake; passing nil causes them to be created from conn.
// isClient reports whether this side is the WebSocket client (which must mask
// its outgoing frames).
func newWSConn(conn net.Conn, br *bufio.Reader, bw *bufio.Writer, isClient bool) *wsConn {
	if br == nil {
		br = bufio.NewReader(conn)
	}
	if bw == nil {
		bw = bufio.NewWriter(conn)
	}
	return &wsConn{conn: conn, br: br, bw: bw, isClient: isClient}
}

// Close closes the underlying connection exactly once.
func (c *wsConn) Close() error {
	var err error
	c.closeOnce.Do(func() { err = c.conn.Close() })
	return err
}

// WriteMessage writes a single unmasked/known-masked data or control frame.
// The caller must hold no locks.
func (c *wsConn) WriteMessage(opcode byte, payload []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()

	b0 := byte(0x80) | opcode // FIN set on the single frame
	if err := c.bw.WriteByte(b0); err != nil {
		return err
	}

	maskBit := byte(0)
	if c.isClient {
		maskBit = 0x80
	}
	if err := writeFrameHeader(c.bw, maskBit, len(payload)); err != nil {
		return err
	}

	var maskKey [4]byte
	if c.isClient {
		if _, err := rand.Read(maskKey[:]); err != nil {
			return err
		}
		if _, err := c.bw.Write(maskKey[:]); err != nil {
			return err
		}
	}

	if !c.isClient {
		if _, err := c.bw.Write(payload); err != nil {
			return err
		}
		return c.bw.Flush()
	}
	if err := writeMaskedPayload(c.bw, payload, maskKey); err != nil {
		return err
	}
	return c.bw.Flush()
}

func writeFrameHeader(w *bufio.Writer, maskBit byte, n int) error {
	switch {
	case n < 126:
		return w.WriteByte(maskBit | byte(n))
	case n <= 0xFFFF:
		if err := w.WriteByte(maskBit | 126); err != nil {
			return err
		}
		var ext [2]byte
		binary.BigEndian.PutUint16(ext[:], uint16(n))
		_, err := w.Write(ext[:])
		return err
	default:
		if err := w.WriteByte(maskBit | 127); err != nil {
			return err
		}
		var ext [8]byte
		binary.BigEndian.PutUint64(ext[:], uint64(n))
		_, err := w.Write(ext[:])
		return err
	}
}

func writeMaskedPayload(w *bufio.Writer, payload []byte, maskKey [4]byte) error {
	var buf [4096]byte
	for off := 0; off < len(payload); off += len(buf) {
		end := off + len(buf)
		if end > len(payload) {
			end = len(payload)
		}
		for i, b := range payload[off:end] {
			buf[i] = b ^ maskKey[(off+i)&3]
		}
		if _, err := w.Write(buf[:end-off]); err != nil {
			return err
		}
	}
	return nil
}

// ReadMessage reads the next complete data message, reassembling fragmented
// frames, replying to pings, and terminating (io.EOF) on a close frame or
// network error. Only one goroutine may call it.
func (c *wsConn) ReadMessage() (byte, []byte, error) {
	var (
		payload []byte
		dataOp  byte
		started bool
	)

	for {
		fin, opcode, data, err := c.readFrame()
		if err != nil {
			return 0, nil, err
		}

		switch opcode {
		case opPing:
			if err := c.WriteMessage(opPong, data); err != nil {
				return 0, nil, err
			}
			continue
		case opPong:
			continue
		case opClose:
			_ = c.WriteMessage(opClose, nil)
			return opClose, nil, io.EOF
		case opContinuation:
			if !started {
				return 0, nil, errors.New("websocket: unexpected continuation frame")
			}
			payload = append(payload, data...)
			if fin {
				return dataOp, payload, nil
			}
		case opText, opBinary:
			if started {
				return 0, nil, errors.New("websocket: new data frame during fragmented message")
			}
			dataOp = opcode
			if fin {
				return opcode, data, nil
			}
			started = true
			payload = append([]byte(nil), data...)
		default:
			return 0, nil, fmt.Errorf("websocket: unsupported opcode %d", opcode)
		}
	}
}

// readFrame reads a single frame header and payload, unmasking when required.
func (c *wsConn) readFrame() (fin bool, opcode byte, payload []byte, err error) {
	var hdr [2]byte
	if _, err = io.ReadFull(c.br, hdr[:]); err != nil {
		return
	}
	fin = hdr[0]&0x80 != 0
	opcode = 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
		}
		length = uint64(binary.BigEndian.Uint16(ext[:]))
	case 127:
		var ext [8]byte
		if _, err = io.ReadFull(c.br, ext[:]); err != nil {
			return
		}
		length = binary.BigEndian.Uint64(ext[:])
	}
	if length > maxFramePayload {
		err = fmt.Errorf("websocket: frame too large (%d bytes)", length)
		return
	}

	var maskKey [4]byte
	if masked {
		if _, err = io.ReadFull(c.br, maskKey[:]); err != nil {
			return
		}
	}

	payload = make([]byte, length)
	if _, err = io.ReadFull(c.br, payload); err != nil {
		return
	}
	if masked {
		for i := range payload {
			payload[i] ^= maskKey[i&3]
		}
	}
	return fin, opcode, payload, nil
}

// computeAccept derives the Sec-WebSocket-Accept value for a handshake key.
func computeAccept(key string) string {
	h := sha1.New()
	_, _ = io.WriteString(h, key+wsGUID)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

// randomWSKey returns a random WebSocket handshake key.
func randomWSKey() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	return base64.StdEncoding.EncodeToString(b)
}

// upgradeWS performs the server side of the WebSocket handshake by hijacking
// the HTTP connection and returning the upgraded wsConn.
func upgradeWS(w http.ResponseWriter, r *http.Request) (*wsConn, error) {
	if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") ||
		!headerHasToken(r.Header, "Connection", "upgrade") {
		http.Error(w, "not a websocket upgrade", http.StatusBadRequest)
		return nil, errors.New("websocket: not an upgrade request")
	}
	key := r.Header.Get("Sec-WebSocket-Key")
	if key == "" {
		http.Error(w, "missing Sec-WebSocket-Key", http.StatusBadRequest)
		return nil, errors.New("websocket: missing handshake key")
	}

	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijacking not supported", http.StatusInternalServerError)
		return nil, errors.New("websocket: hijacking not supported")
	}
	netConn, brw, err := hj.Hijack()
	if err != nil {
		return nil, fmt.Errorf("websocket: hijack: %w", err)
	}

	resp := "HTTP/1.1 101 Switching Protocols\r\n" +
		"Upgrade: websocket\r\n" +
		"Connection: Upgrade\r\n" +
		"Sec-WebSocket-Accept: " + computeAccept(key) + "\r\n\r\n"
	if _, err := brw.WriteString(resp); err != nil {
		_ = netConn.Close()
		return nil, err
	}
	if err := brw.Flush(); err != nil {
		_ = netConn.Close()
		return nil, err
	}
	return newWSConn(netConn, brw.Reader, brw.Writer, false), nil
}

// headerHasToken reports whether a comma-separated header value contains the
// given token case-insensitively.
func headerHasToken(h http.Header, name, token string) bool {
	for _, val := range h.Values(name) {
		for _, part := range strings.Split(val, ",") {
			if strings.EqualFold(strings.TrimSpace(part), token) {
				return true
			}
		}
	}
	return false
}

// dialWS performs the client side of the WebSocket handshake, returning the
// upgraded wsConn. accessToken, when non-empty, is sent both as an
// Authorization: Bearer header and as an access_token query parameter for
// OneBot implementations that expect either.
func dialWS(ctx context.Context, dial DialContextFunc, rawURL, accessToken string) (*wsConn, error) {
	if dial == nil {
		dial = (&net.Dialer{Timeout: 15 * time.Second}).DialContext
	}
	u, err := url.Parse(rawURL)
	if err != nil {
		return nil, fmt.Errorf("websocket: parse url %q: %w", rawURL, err)
	}
	if u.Scheme != "ws" && u.Scheme != "wss" {
		return nil, fmt.Errorf("websocket: unsupported scheme %q (want ws/wss)", u.Scheme)
	}

	host := u.Host
	if u.Port() == "" {
		if u.Scheme == "wss" {
			host += ":443"
		} else {
			host += ":80"
		}
	}
	conn, err := dial(ctx, "tcp", host)
	if err != nil {
		return nil, fmt.Errorf("websocket: dial %s: %w", host, err)
	}

	if u.Scheme == "wss" {
		tlsConn := tls.Client(conn, &tls.Config{ServerName: u.Hostname()})
		if err := tlsConn.HandshakeContext(ctx); err != nil {
			_ = conn.Close()
			return nil, fmt.Errorf("websocket: tls handshake: %w", err)
		}
		conn = tlsConn
	}

	brw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
	key := randomWSKey()
	var b strings.Builder
	b.WriteString("GET " + u.RequestURI() + " HTTP/1.1\r\n")
	b.WriteString("Host: " + u.Host + "\r\n")
	b.WriteString("Upgrade: websocket\r\n")
	b.WriteString("Connection: Upgrade\r\n")
	b.WriteString("Sec-WebSocket-Key: " + key + "\r\n")
	b.WriteString("Sec-WebSocket-Version: 13\r\n")
	if accessToken != "" {
		b.WriteString("Authorization: Bearer " + accessToken + "\r\n")
	}
	b.WriteString("\r\n")
	if _, err := brw.WriteString(b.String()); err != nil {
		_ = conn.Close()
		return nil, err
	}
	if err := brw.Flush(); err != nil {
		_ = conn.Close()
		return nil, err
	}

	resp, err := http.ReadResponse(brw.Reader, &http.Request{Method: "GET"})
	if err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("websocket: read handshake response: %w", err)
	}
	if resp.StatusCode != http.StatusSwitchingProtocols {
		_ = conn.Close()
		return nil, fmt.Errorf("websocket: unexpected handshake status %d", resp.StatusCode)
	}
	if got := resp.Header.Get("Sec-WebSocket-Accept"); got != computeAccept(key) {
		_ = conn.Close()
		return nil, errors.New("websocket: bad Sec-WebSocket-Accept in response")
	}
	return newWSConn(conn, brw.Reader, brw.Writer, true), nil
}