package main

import (
	"bufio"
	"fmt"
	"log"
	"sync"
)

type PacketQueue struct {
	mu               sync.RWMutex
	expectedPackages map[uint32]*Listener
}

type CallResult struct {
	StatusCode uint32
	Data       []byte
}

type Listener struct {
	Count int
	Chan  chan CallResult
}

func NewPacketQueue(connection *PersistentConnection) *PacketQueue {
	queue := &PacketQueue{
		expectedPackages: make(map[uint32]*Listener),
	}
	go queue.HandleConnection(connection)
	return queue
}

func (p *PacketQueue) RemoveListeners() {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, l := range p.expectedPackages {
		close(l.Chan)
	}
	p.expectedPackages = make(map[uint32]*Listener)
}

func (p *PacketQueue) HandleConnection(connection *PersistentConnection) error {
	defer connection.Close()
	defer p.RemoveListeners()
	var packet Packet
	reader := bufio.NewReader(connection)

	for {
		err := ReadPacket(reader, &packet)
		if err != nil {
			return connection.HandleConnectionError(err)
		}
		if packet.Version != CurrentPacketVersion {
			log.Printf("Incorrect packet version: %v\n", packet.Version)
			return connection.HandleConnectionError(fmt.Errorf("incorrect packet version: %d", packet.Version))
		}
		if packet.DataLength == 0 {
			go p.HandleData(packet.MessageType, CallResult{
				StatusCode: packet.StatusCode,
				Data:       []byte{},
			})
			continue
		}
		data, err := GetPacketData(reader, packet.DataLength)
		if err != nil {
			log.Printf("Error receiving packet data: %v\n", err)
			return connection.HandleConnectionError(err)
		} else {
			go p.HandleData(packet.MessageType, CallResult{
				StatusCode: packet.StatusCode,
				Data:       data,
			})
		}
	}
}

func (p *PacketQueue) HandleData(t uint32, data CallResult) {
	p.mu.Lock()
	defer p.mu.Unlock()
	l, ok := p.expectedPackages[t]
	if ok {
		l.Chan <- data
		l.Count--
		if l.Count == 0 {
			close(l.Chan)
			delete(p.expectedPackages, t)
		}
		return
	}
}

func (p *PacketQueue) Expect(messageType uint32) <-chan CallResult {
	p.mu.Lock()
	defer p.mu.Unlock()
	l, ok := p.expectedPackages[messageType]
	if ok {
		l.Count++
		return l.Chan
	}

	ch := make(chan CallResult)
	p.expectedPackages[messageType] = &Listener{
		Count: 1,
		Chan:  ch,
	}

	return ch

}