package hopp

import "io"
import "fmt"
import "net"
import "sync"
import "git.tebibyte.media/sashakoshka/hopp/tape"
import "git.tebibyte.media/sashakoshka/go-util/sync"

const closeMethod = 0xFFFF
const int64Max = int64((^uint64(0)) >> 1)

// Party represents a side of a connection.
type Party bool; const (
	ServerSide Party = false
	ClientSide Party = true
)

type a struct {
	sizeLimit    int
	underlying   net.Conn
	party        Party
	transID      int64
	transLock    sync.RWMutex
	sendLock     sync.Mutex
	transMap     map[int64] *transA
	transChan    chan *transA
	done         chan struct { }
	err          error
}

// AdaptA returns a connection implementing METADAPT-A over a singular stream-
// oriented transport such as TCP or UNIX domain stream sockets.
func AdaptA(underlying net.Conn, party Party) Conn {
	conn := &a {
		underlying: underlying,
		party:      party,
		transMap:   make(map[int64] *transA),
		transChan:  make(chan *transA),
		done:       make(chan struct { }),
	}
	if party == ClientSide {
		conn.transID =  1
	} else {
		conn.transID = -1
	}
	go conn.receive()
	return conn
}

func (this *a) Close() error {
	close(this.done)
	return this.underlying.Close()
}

func (this *a) LocalAddr() net.Addr {
	return this.underlying.LocalAddr()
}

func (this *a) RemoteAddr() net.Addr {
	return this.underlying.RemoteAddr()
}

func (this *a) OpenTrans() (Trans, error) {
	this.transLock.Lock()
	defer this.transLock.Unlock()
	id := this.transID
	this.transID ++
	trans := &transA {
		parent:   this,
		id:       id,
		incoming: usync.NewGate[incomingMessage](),
	}
	this.transMap[id] = trans
	if this.transID == int64Max {
		return nil, fmt.Errorf("could not open transaction: %w", ErrIntegerOverflow)
	}
	this.transID ++
	return trans, nil
}

func (this *a) AcceptTrans() (Trans, error) {
	select {
	case trans := <- this.transChan:
		return trans, nil
	case <- this.done:
		return nil, fmt.Errorf("could not accept transaction: %w", io.EOF)
	}
}

func (this *a) SetSizeLimit(limit int) {
	this.sizeLimit = limit
}

func (this *a) unlistTransactionSafe(id int64) {
	this.transLock.Lock()
	defer this.transLock.Unlock()
	delete(this.transMap, id)
}

func (this *a) sendMessageSafe(trans int64, method uint16, data []byte) error {
	this.sendLock.Lock()
	defer this.sendLock.Unlock()
	return encodeMessageA(this.underlying, trans, method, data)
}

func (this *a) receive() {
	defer func() {
		this.underlying.Close()
		this.transLock.Lock()
		defer this.transLock.Unlock()
		for _, trans := range this.transMap {
			trans.closeDontUnlist()
		}
		clear(this.transMap)
	}()
	for {
		transID, method, chunked, payload, err := decodeMessageA(this.underlying)
		if err != nil {
			this.err = fmt.Errorf("could not receive message: %w", err)
			return
		}

		err = this.receiveMultiplex(transID, method, chunked, payload)
		if err != nil {
			this.err = fmt.Errorf("could not receive message: %w", err)
			return
		}
	}
}

func (this *a) receiveMultiplex(transID int64, method uint16, chunked bool, payload []byte) error {
	if transID == 0 { return ErrMessageMalformed }

	trans, err := func() (*transA, error) {
		this.transLock.Lock()
		defer this.transLock.Unlock()

		trans, ok := this.transMap[transID]
		if !ok {
			// it is forbidden for the other party to initiate a transaction
			// with an ID from this party
			if this.party == partyFromTransID(transID) {
				return nil, ErrMessageMalformed
			}
			trans = &transA {
				parent:   this,
				id:       transID,
				incoming: usync.NewGate[incomingMessage](),
			}
			this.transMap[transID] = trans
			this.transChan <- trans
		}
		return trans, nil
	}()
	if err != nil { return err }

	trans.incoming.Send(incomingMessage {
		method:  method,
		chunked: chunked,
		payload: payload,
	})
	return nil
}

type transA struct {
	parent        *a
	id            int64
	incoming      usync.Gate[incomingMessage]
	currentReader io.Reader
}

func (this *transA) Close() error {
	err := this.closeDontUnlist()
	this.parent.unlistTransactionSafe(this.ID())
	return err
}

func (this *transA) closeDontUnlist() error {
	this.Send(closeMethod, nil)
	return this.incoming.Close()
}

func (this *transA) ID() int64 {
	return this.id
}

func (this *transA) Send(method uint16, data []byte) error {
	return this.parent.sendMessageSafe(this.id, method, data)
}

func (this *transA) Receive() (method uint16, data []byte, err error) {
	method, reader, err := this.ReceiveReader()
	if err != nil { return 0, nil, err }
	data, err = io.ReadAll(reader)
	if err != nil { return 0, nil, err }
	return method, data, nil
}

func (this *transA) ReceiveReader() (uint16, io.Reader, error) {
	// drain previous reader if necessary
	if this.currentReader != nil {
		io.Copy(io.Discard, this.currentReader)
	}

	// create new reader
	reader := &readerA {
		parent: this,
	}
	method, err := reader.pull()
	if err != nil { return 0, nil, err}
	this.currentReader = reader
	return method, reader, nil
}

type readerA struct {
	parent   *transA
	leftover []byte
	eof      bool
}

func (this *readerA) pull() (uint16, error) {
	// if the previous message ended the chain, return an io.EOF
	if this.eof {
		return 0, io.EOF
	}

	// get a message from the transaction we are a part of
	receive := this.parent.incoming.Receive()
	if receive != nil {
		if message, ok := <- receive; ok {
			if message.method != closeMethod {
				this.leftover = append(this.leftover, message.payload...)
				if !message.chunked {
					this.eof = true
				}
			}
		}
	}

	// close and return error on failure
	this.eof = true
	this.parent.Close()
	if this.parent.parent.err == nil {
		return 0, fmt.Errorf("could not receive message: %w", io.EOF)
	} else {
		return 0, this.parent.parent.err
	}
}

func (this *readerA) Read(buffer []byte) (int, error) {
	if len(this.leftover) == 0 {
		if this.eof { return 0, io.EOF }
		this.pull()
	}

	copied := copy(buffer, this.leftover)
	this.leftover = this.leftover[copied:]
	return copied, nil
}

type incomingMessage struct {
	method  uint16
	chunked bool
	payload []byte
}

func encodeMessageA(writer io.Writer, trans int64, method uint16, data []byte) error {
	buffer := make([]byte, 12 + len(data))
	tape.EncodeI64(buffer[:8], trans)
	tape.EncodeI16(buffer[8:10], method)
	length, ok := tape.U16CastSafe(len(data))
	if !ok { return ErrPayloadTooLarge }
	tape.EncodeI16(buffer[10:12], length)
	copy(buffer[12:], data)
	_, err := writer.Write(buffer)
	return err
}

func decodeMessageA(reader io.Reader) (int64, uint16, bool, []byte, error) {
	headerBuffer := [18]byte { }
	_, err := io.ReadFull(reader, headerBuffer[:])
	if err != nil { return 0, 0, false, nil, err }
	transID, err := tape.DecodeI64[int64](headerBuffer[:8])
	if err != nil { return 0, 0, false, nil, err }
	method, err := tape.DecodeI16[uint16](headerBuffer[8:10])
	if err != nil { return 0, 0, false, nil, err }
	size, err := tape.DecodeI64[uint64](headerBuffer[10:18])
	if err != nil { return 0, 0, false, nil, err }
	chunked, size := splitCCBSize(size)
	payloadBuffer := make([]byte, int(size))
	_, err = io.ReadFull(reader, payloadBuffer)
	if err != nil { return 0, 0, false, nil, err }
	return transID, method, chunked, payloadBuffer, nil
}

func partyFromTransID(id int64) Party {
	return id > 0
}

func splitCCBSize(size uint64) (bool, uint64) {
	return size >> 63 > 1, size & 0x7FFFFFFFFFFFFFFF
}