package transport

import (
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/hex"
	"syscall"
)

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) (byte, []byte, int) {
	if len(data) < 2 {
		return 0, nil, 0
	}
	op := data[0] & 0x0f
	masked := data[1]&0x80 != 0
	length := int(data[1] & 0x7f)
	pos := 2

	if length == 126 {
		if len(data) < 4 {
			return 0, nil, 0
		}
		length = int(data[2])<<8 | int(data[3])
		pos = 4
	} else if length == 127 {
		if len(data) < 10 {
			return 0, nil, 0
		}
		length = int(data[6])<<24 | int(data[7])<<16 | int(data[8])<<8 | int(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) []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 = append(buf, hdr[:n]...)
	buf = append(buf, payload...)
	return buf
}

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

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

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

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

func computeAccept(key string) 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)
	}
}

func (s *Server) processWS(c *tconn) {
	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:
		s.handler.OnWSMessage(c.fd, payload)
	}
}

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