package main

import (
	"bytes"
	"context"
	"fmt"
	"net"
	"strconv"
	"sync"
	"time"
)

type ClientConn struct {
	conn *net.TCPConn
}

type Client struct {
	address string

	mu sync.Mutex
	wg sync.WaitGroup

	closed  chan struct{}
	waiters []chan error

	pool [4]*ClientConn
	len  uint8

	queue chan []byte
}

func NewClient(address string) Client {
	return Client{
		address: address,
		queue:   make(chan []byte, 1024),
	}
}

func (c *Client) Close() {
	c.close()
	c.wg.Wait()
	close(c.queue)
}

func (c *Client) close() {
	c.mu.Lock()
	defer c.mu.Unlock()

	for _, waiter := range c.waiters {
		waiter <- context.Canceled
		close(waiter)
	}
	c.waiters = c.waiters[:0]

	c.closed = make(chan struct{})
	close(c.closed)

	for _, conn := range c.pool[:c.len] {
		if conn.conn != nil {
			conn.conn.Close()
		}
	}
}

// spawnOrWaitForAvailableConnection returns a waiter channel if we are waiting for a connection
// in our pool (which may or may not have just been spawned) to be ready, or nil otherwise. It
// executes a pooling strategy.
//
// if any connected, pool isn't full, write queue isn't empty, spawn a new connection
// if any connected, pool isn't full, write queue is empty, don't spawn a new connection
// if any connected, pool is full, don't spawn a new connection

// if none connected, pool isn't full, write queue isn't empty, spawn a new connection
// if none connected, pool isn't full, write queue is empty, spawn a new connection
// if none connected, pool is full, don't spawn a new connection
func (c *Client) spawnOrWaitForAvailableConnection() (chan error, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed != nil {
		return nil, context.Canceled
	}

	pool := c.pool[0:c.len:cap(c.pool)]

	connected := false
	for _, conn := range pool {
		if conn.conn != nil {
			connected = true
			break
		}
	}

	if len(pool) < cap(c.pool) && (!connected || len(c.queue) > 0) {
		cc := &ClientConn{}
		c.pool[c.len] = cc
		c.len = c.len + 1
		c.wg.Add(1)

		go func() {
			defer c.wg.Done()

			for {
				loop := func() bool {
					var d net.Dialer

					conn, err := func() (net.Conn, error) {
						ctx, cancel := context.WithTimeout(context.Background(), 5 * time.Second)
						defer cancel()

						return d.DialContext(ctx, "tcp", c.address)
					}()
					if err != nil {
						if !c.reportConnectionError(cc, err) {
							return false
						}
						return true
					}

					if !c.reportConnected(cc, conn.(*net.TCPConn)) {
						conn.Close()
						return false
					}

					writerDone := make(chan error)
					writerStop := make(chan struct{})
					readerDone := make(chan error)

					go func() {
						writerDone <- func() error {
							for {
								select {
								case <-writerStop:
									return nil
								case buf, open := <-c.queue:
									if !open {
										return context.Canceled
									}

									_, err := cc.conn.Write(buf)
									if err != nil {
										return err
									}
								}
							}
						}()
					}()

					go func() {
						readerDone <- func() error {
							var buf [1024]byte
							for {
								n, err := cc.conn.Read(buf[0:])
								if err != nil {
									return err
								}
								if n == 0 {
									return nil
								}
								fmt.Printf("Got: '%s'\n", bytes.TrimSpace(buf[0:n]))
							}
						}()
					}()

					select {
					case <-writerDone:
						conn.Close()
						<-readerDone
					case <-readerDone:
						conn.Close()
						close(writerStop)
						<-writerDone
					}

					return !c.reportDisconnected(cc)
				}()

				if !loop {
					break
				}
			}
		}()
	}

	if !connected {
		waiter := make(chan error, 1)
		c.waiters = append(c.waiters, waiter)
		return waiter, nil
	}

	return nil, nil
}

func (c *Client) tryRemoveWaiter(waiter chan error) bool {
	// Attempt to remove our waiter entry from the waiters list. It is possible
	// that our waiter has already been removed from the waiter list.

	// If we did not find our waiter entry, then that means that a connection is available.
	// This is needed because we are racing whether or not a pool connection is available
	// against ctx.Done().

	c.mu.Lock()
	defer c.mu.Unlock()

	waiters := c.waiters[:0]
	found := false

	for _, w := range c.waiters {
		if w == waiter {
			found = true
			continue
		}
		waiters = append(waiters, w)
	}
	c.waiters = waiters

	if !found {
		return false
	}

	close(waiter)
	return true
}

func (c *Client) reportConnected(cc *ClientConn, dialed *net.TCPConn) bool {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed != nil {
		pool := c.pool[:0]
		for _, conn := range c.pool {
			if cc == conn {
				continue
			}
			pool = append(pool, conn)
		}
		c.len = c.len - 1

		return false
	}

	cc.conn = dialed

	for _, waiter := range c.waiters {
		close(waiter)
	}
	c.waiters = c.waiters[:0]

	return true
}

func (c *Client) reportConnectionError(cc *ClientConn, err error) bool {
	c.mu.Lock()
	defer c.mu.Unlock()

	// If cc is not the last connection left, do not report
	// the error to the existing waiters.

	// Remove client connection entry from the pool.

	if c.len > 1 {
		pool := c.pool[:0]
		for _, conn := range c.pool {
			if cc == conn {
				continue
			}
			pool = append(pool, conn)
		}
		c.len = c.len - 1

		return false
	}

	for _, waiter := range c.waiters {
		waiter <- err
		close(waiter)
	}
	c.waiters = c.waiters[:0]

	return true
}

func (c *Client) reportDisconnected(cc *ClientConn) bool {
	c.mu.Lock()
	defer c.mu.Unlock()

	cc.conn = nil

	if c.closed != nil || c.len > 1 {
		pool := c.pool[:0]
		for _, conn := range c.pool {
			if cc == conn {
				continue
			}
			pool = append(pool, conn)
		}
		c.len = c.len - 1

		return true
	}

	return false
}

func (c *Client) EnsureConnectionAvailable(ctx context.Context) error {
	waiter, err := c.spawnOrWaitForAvailableConnection()
	if err != nil {
		return err
	}
	if waiter == nil {
		return nil
	}

	select {
	case <-c.closed:
		if !c.tryRemoveWaiter(waiter) {
			return nil
		}
		return context.Canceled
	case <-ctx.Done():
		if !c.tryRemoveWaiter(waiter) {
			return nil
		}
		return ctx.Err()
	case err := <-waiter:
		return err
	}
}

func (c *Client) Write(ctx context.Context, buf []byte) error {
	if err := c.EnsureConnectionAvailable(ctx); err != nil {
		return err
	}

	select {
	case <-c.closed:
		return context.Canceled
	case <-ctx.Done():
		return ctx.Err()
	case c.queue <- buf:
		return nil
	}
}

func main() {
	client := NewClient("127.0.0.1:9000")
	defer client.Close()

	for i := 0; i < 100; i++ {
		func() {
			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()

			if err := client.Write(ctx, append([]byte(strconv.FormatInt(int64(i), 10)), '\n')); err != nil {
				panic(err)
			}
		}()
	}
}