package wire

import (
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
)

const (
	RequestHeaderSize  = 29
	ResponseHeaderSize = 5
	MaxPayload         = 2 * 1024 * 1024

	ModeProbe    byte = 0
	ModeOpen     byte = 1
	ModeUpload   byte = 2
	ModeDownload byte = 3
	ModeClose    byte = 4

	StatusOK    byte = 0
	StatusError byte = 1
	StatusData  byte = 2
	StatusWait  byte = 3
	StatusEOF   byte = 4

	ProbeUpload    byte = 1
	ProbeDownload  byte = 2
	ProbeKeepalive byte = 3
	ProbeBatch     byte = 4
)

var ProbeMagic = [4]byte{'D', 'T', 'P', '2'}

type SessionID [16]byte

type Request struct {
	Mode    byte
	Session SessionID
	Seq     uint64
	Payload []byte
}

func MaskInPlace(data []byte, sid SessionID, mode byte, seq uint64, response bool) {
	if len(data) == 0 {
		return
	}

	var seed [30]byte
	copy(seed[:16], sid[:])
	seed[16] = mode
	binary.BigEndian.PutUint64(seed[17:25], seq)
	if response {
		seed[25] = 1
	}

	var counter uint32
	for off := 0; off < len(data); {
		binary.BigEndian.PutUint32(seed[26:30], counter)
		block := sha256.Sum256(seed[:])
		n := len(data) - off
		if n > len(block) {
			n = len(block)
		}
		for i := 0; i < n; i++ {
			data[off+i] ^= block[i]
		}
		off += n
		counter++
	}
}

func WriteRequest(w io.Writer, mode byte, sid SessionID, seq uint64, plaintext []byte) error {
	if len(plaintext) > MaxPayload {
		return fmt.Errorf("request payload too large: %d", len(plaintext))
	}

	packet := make([]byte, RequestHeaderSize+len(plaintext))
	packet[0] = mode
	copy(packet[1:17], sid[:])
	binary.BigEndian.PutUint64(packet[17:25], seq)
	binary.BigEndian.PutUint32(packet[25:29], uint32(len(plaintext)))
	copy(packet[29:], plaintext)
	MaskInPlace(packet[29:], sid, mode, seq, false)
	return writeAll(w, packet)
}

func ReadRequest(r io.Reader) (Request, error) {
	var req Request
	var header [RequestHeaderSize]byte
	if _, err := io.ReadFull(r, header[:]); err != nil {
		return req, err
	}

	req.Mode = header[0]
	copy(req.Session[:], header[1:17])
	req.Seq = binary.BigEndian.Uint64(header[17:25])
	n := binary.BigEndian.Uint32(header[25:29])
	if n > MaxPayload {
		return req, errors.New("request payload too large")
	}

	if n > 0 {
		req.Payload = make([]byte, int(n))
		if _, err := io.ReadFull(r, req.Payload); err != nil {
			return req, err
		}
		MaskInPlace(req.Payload, req.Session, req.Mode, req.Seq, false)
	}
	return req, nil
}

func WriteResponse(w io.Writer, status byte, body []byte) error {
	if len(body) > MaxPayload {
		return fmt.Errorf("response body too large: %d", len(body))
	}
	packet := make([]byte, ResponseHeaderSize+len(body))
	packet[0] = status
	binary.BigEndian.PutUint32(packet[1:5], uint32(len(body)))
	copy(packet[5:], body)
	return writeAll(w, packet)
}

func WriteMaskedResponse(w io.Writer, status byte, body []byte, sid SessionID, mode byte, seq uint64) error {
	if len(body) > MaxPayload {
		return fmt.Errorf("response body too large: %d", len(body))
	}
	packet := make([]byte, ResponseHeaderSize+len(body))
	packet[0] = status
	binary.BigEndian.PutUint32(packet[1:5], uint32(len(body)))
	copy(packet[5:], body)
	MaskInPlace(packet[5:], sid, mode, seq, true)
	return writeAll(w, packet)
}

func ReadResponse(r io.Reader) (byte, []byte, error) {
	var header [ResponseHeaderSize]byte
	if _, err := io.ReadFull(r, header[:]); err != nil {
		return 0, nil, err
	}
	n := binary.BigEndian.Uint32(header[1:5])
	if n > MaxPayload {
		return 0, nil, errors.New("response body too large")
	}
	var body []byte
	if n > 0 {
		body = make([]byte, int(n))
		if _, err := io.ReadFull(r, body); err != nil {
			return 0, nil, err
		}
	}
	return header[0], body, nil
}

func DecodeMaskedResponse(status byte, body []byte, sid SessionID, mode byte, seq uint64) []byte {
	if len(body) == 0 || status == StatusError {
		return body
	}
	out := append([]byte(nil), body...)
	MaskInPlace(out, sid, mode, seq, true)
	return out
}

func writeAll(w io.Writer, b []byte) error {
	for len(b) > 0 {
		n, err := w.Write(b)
		if err != nil {
			return err
		}
		if n <= 0 {
			return io.ErrShortWrite
		}
		b = b[n:]
	}
	return nil
}