package transport

import (
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"syscall"
	"git.smesh.lol/morly/pkg/metrics"
)

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

// parseWSFrame extracts one WS frame from data.
// Returns (op, payload, consumed). consumed=0 means incomplete.
func parseWSFrame(data []byte) (op byte, payload []byte, n int32) {
	if len(data) < 2 {
		return 0, nil, 0
	}
	op := data[0] & 0x0f
	masked := data[1]&0x80 != 0
	length := int32(data[1] & 0x7f)
	pos := 2

	if length == 126 {
		if len(data) < 4 {
			return 0, nil, 0
		}
		length = int32(data[2])<<8 | int32(data[3])
		pos = 4
	} else if length == 127 {
		if len(data) < 10 {
			return 0, nil, 0
		}
		length = int32(data[6])<<24 | int32(data[7])<<16 | int32(data[8])<<8 | int32(data[9])
		pos = 10
	}

	var mask [4]byte
	if masked {
		if len(data) < pos+4 {
			return 0, nil, 0
		}
		copy(mask[:], data[pos:pos+4])
		pos += 4
	}
	if len(data) < pos+length {
		return 0, nil, 0
	}

	payload := makeCopy(data[pos : pos+length])
	if masked {
		for i := range payload {
			payload[i] ^= mask[i%4]
		}
	}
	return op, payload, pos + length
}

func buildWSFrame(op byte, payload []byte) (buf []byte) {
	plen := len(payload)
	var hdr [10]byte
	hdr[0] = 0x80 | op
	n := 2
	if plen < 126 {
		hdr[1] = byte(plen)
	} else if plen < 65536 {
		hdr[1] = 126
		hdr[2] = byte(plen >> 8)
		hdr[3] = byte(plen)
		n = 4
	} else {
		hdr[1] = 127
		hdr[6] = byte(plen >> 24)
		hdr[7] = byte(plen >> 16)
		hdr[8] = byte(plen >> 8)
		hdr[9] = byte(plen)
		n = 10
	}
	buf := []byte{:0:n + plen}
	buf = buf | hdr[:n]
	buf = buf | payload
	return buf
}

func writeWSFrame(fd int32, op byte, payload []byte) {
	writeAll(fd, buildWSFrame(op, payload))
}

func writeWSFrameErr(fd int32, op byte, payload []byte) (eerr error) {
	return writeAll(fd, buildWSFrame(op, payload))
}

func writeWSClose(fd int32) {
	writeWSFrame(fd, opClose, []byte{0x03, 0xe8})
}

func writeAll(fd int32, data []byte) (eerr error) {
	for len(data) > 0 {
		n, err := syscall.Write(fd, data)
		if n > 0 {
			data = data[n:]
		}
		if errAgain(err) {
			return syscall.EAGAIN
		}
		if errIntr(err) {
			continue
		}
		if err != nil {
			return err
		}
	}
	return nil
}

func computeAccept(key string) (s string) {
	h := sha1.New()
	h.Write([]byte(key))
	h.Write([]byte(wsMagic))
	return base64.StdEncoding().EncodeToString(h.Sum(nil))
}

func (s *Server) upgradeWS(c *tconn, req *httpReq) {
	key := req.headers["sec-websocket-key"]
	if key == "" {
		writeHTTPResponse(c.fd, 400, nil, []byte("missing Sec-WebSocket-Key"))
		s.closeConn(c)
		return
	}
	wl, allow := s.handler.OnWSUpgrade(c.fd, c.remoteIP, s.ipConns[c.remoteIP])
	if !allow {
		writeHTTPResponse(c.fd, 429, nil, []byte("too many websockets"))
		s.closeConn(c)
		return
	}
	s.ipConns[c.remoteIP]++
	c.whitelisted = wl
	accept := computeAccept(key)
	resp := "HTTP/1.1 101 Switching Protocols\r\n" |
		"Upgrade: websocket\r\n" |
		"Connection: Upgrade\r\n" |
		"Sec-WebSocket-Accept: " | accept | "\r\n" |
		"\r\n"
	writeAll(c.fd, []byte(resp))
	c.phase = phaseWS
	s.handler.OnWSConnected(c.fd)
	if c.wpos > 0 {
		s.processWS(c)
	}
}

// processWS drains every complete frame already in the connection's buffer.
//
// Handling one frame per call is not enough: readConn reads whatever the socket
// has, and a client is free to put several frames in one segment - a browser
// sending four REQs back to back does exactly that. The frames after the first
// are then already out of the socket, so no further EPOLLIN ever arrives for
// them and they sat in c.buf unprocessed. A pipelined client got one answer and
// every later subscription on that connection was silently dead, which is the
// whole of the browser-feed failure: four requests, one EOSE, no stored event.
func (s *Server) processWS(c *tconn) {
	var op byte
	var payload []byte
	var consumed int32
	var frameStart int64
	for {
		op, payload, consumed = parseWSFrame(c.buf[:c.wpos])
		if consumed == 0 {
			return
		}
		copy(c.buf, c.buf[consumed:c.wpos])
		c.wpos -= consumed

		switch op {
		case opClose:
			writeWSClose(c.fd)
			s.closeConn(c)
			return
		case opPing:
			writeWSFrame(c.fd, opPong, payload)
		case opText, opBin:
			frameStart = metrics.Now()
			s.handler.OnWSMessage(c.fd, payload)
			metrics.WSFrameNs.Observe(metrics.Since(frameStart))
		}
		// The handler may have closed the connection.
		if s.conns[c.fd] == nil {
			return
		}
	}
}

// Challenge generates a 32-byte random hex challenge string for NIP-42 auth.
func Challenge() (buf []byte) {
	var cb [32]byte
	rand.Read(cb[:])
	out := []byte{:64}
	hex.Encode(out, cb[:])
	return out
}