package src

import (
	"bufio"
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"hash/crc32"
	"io"
	"log"
	"net"
	"os"
	"slices"
	"strings"
	"sync"
)

type Config struct {
	ctx      context.Context
	Listen   string `json:"LISTEN"`
	Forwards string `json:"FORWARDS"`
	Hello    string `json:"HELLO_B64"`
	forwards []*sync.Pool
}

func NewConfig(ctx context.Context) (Config, error) {
	config := Config{
		ctx:      ctx,
		Listen:   ":10000",
		Forwards: "",
		Hello:    base64.StdEncoding.EncodeToString([]byte("*1\r\n$4\r\nping\r\n")),
	}

	if err := config.loadEnv(); err != nil {
		return config, err
	}

	if err := config.setForwards(); err != nil {
		return config, err
	}

	return config, nil
}

func (c *Config) loadEnv() error {
	b, _ := json.Marshal(*c)
	var m map[string]any
	if err := json.Unmarshal(b, &m); err != nil {
		return err
	}
	for k := range m {
		if v := os.Getenv(k); v != "" {
			m[k] = v
		}
	}
	b2, _ := json.Marshal(m)
	return json.Unmarshal(b2, c)
}

func (c *Config) setForwards() error {
	hello, err := base64.StdEncoding.DecodeString(c.Hello)
	if err != nil {
		return err
	}

	addresses := strings.Split(c.Forwards, ",")
	addresses = slices.DeleteFunc(addresses, func(s string) bool { return s == "" })
	if len(addresses) == 0 {
		return fmt.Errorf("at least one $FORWARD required")
	}

	c.forwards = make([]*sync.Pool, len(addresses))
	for i := range addresses {
		address := addresses[i]
		c.forwards[i] = &sync.Pool{
			New: func() any {
				conn, err := c.dial(address, hello)
				if err != nil {
					log.Printf("! failed to dial %q: %v", address, err)
				}
				return conn
			},
		}
	}
	return nil
}

func (c *Config) dial(address string, hello []byte) (net.Conn, error) {
	if err := c.ctx.Err(); err != nil {
		return nil, err
	}

	log.Printf("dialing %q with %q", address, hello)
	v, err := (&net.Dialer{}).DialContext(c.ctx, "tcp", address)
	if err != nil {
		return nil, err
	}

	if _, err := v.Write([]byte(hello)); err != nil {
		v.Close()
		return nil, err
	}

	if raw, err := readMessage(bufio.NewReader(v)); err != nil {
		v.Close()
		return nil, err
	} else if bytes.HasPrefix(raw, []byte("-")) {
		v.Close()
		return nil, fmt.Errorf("failed hello: %s", raw)
	}

	return v, nil
}

func (c Config) Close() {
	for i := range c.forwards {
		if c.forwards[i] != nil {
			i := i
			func() {
				c.forwards[i].New = nil
				for {
					got := c.forwards[i].Get()
					if got == nil {
						break
					}
					got.(net.Conn).Close()
				}
			}()
		}
	}
}

func (c Config) WithConn(hashKey string, foo func(net.Conn) error) error {
	hash := int(crc32.ChecksumIEEE([]byte(hashKey)))
	hashIdx := hash % len(c.forwards)
	forward := c.forwards[hashIdx]

	log.Printf("acquire conn %v", hashIdx)
	forwardCon := forward.Get()
	if forwardCon == nil {
		log.Printf("got a nil conn to %v", hashIdx)
		return io.EOF
	}
	forwardConn := forwardCon.(net.Conn)

	if err := foo(forwardConn); err != nil {
		log.Printf("errored with conn %v: %v", hashIdx, err)
		forwardConn.Close()
		return err
	}

	log.Printf("release conn %v", hashIdx)
	forward.Put(forwardConn)
	return nil
}