package sessionstream

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
)

const (
	ChannelStdin   byte = 0x01
	ChannelStdout  byte = 0x02
	ChannelStderr  byte = 0x03
	ChannelControl byte = 0x04
	FormatV1            = "stdio_mux_v1"
)

type ControlMessage struct {
	Type     string `json:"type"`
	ExitCode *int   `json:"exit_code,omitempty"`
	Error    string `json:"error,omitempty"`
}

func WriteFrame(w io.Writer, channel byte, payload []byte) error {
	var header [5]byte
	header[0] = channel
	binary.BigEndian.PutUint32(header[1:], uint32(len(payload)))
	if _, err := w.Write(header[:]); err != nil {
		return err
	}
	if len(payload) == 0 {
		return nil
	}
	_, err := w.Write(payload)
	return err
}

func ReadFrame(r io.Reader) (byte, []byte, error) {
	var header [5]byte
	if _, err := io.ReadFull(r, header[:]); err != nil {
		return 0, nil, err
	}
	length := binary.BigEndian.Uint32(header[1:])
	payload := make([]byte, length)
	if _, err := io.ReadFull(r, payload); err != nil {
		return 0, nil, err
	}
	return header[0], payload, nil
}

func WriteControl(w io.Writer, message ControlMessage) error {
	payload, err := json.Marshal(message)
	if err != nil {
		return err
	}
	return WriteFrame(w, ChannelControl, payload)
}

func ReadControl(payload []byte) (ControlMessage, error) {
	var message ControlMessage
	if err := json.Unmarshal(payload, &message); err != nil {
		return ControlMessage{}, err
	}
	return message, nil
}

func ReadNextControl(r io.Reader) (ControlMessage, error) {
	channel, payload, err := ReadFrame(r)
	if err != nil {
		return ControlMessage{}, err
	}
	if channel != ChannelControl {
		return ControlMessage{}, fmt.Errorf("unexpected channel %d", channel)
	}
	return ReadControl(payload)
}