package src

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

type Config struct {
	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{
		Listen:   ":10000",
		Forwards: "",
		Hello:    base64.StdEncoding.EncodeToString([]byte("*1\r\n$4\r\nping\r\n")),
	}
	b, _ := json.Marshal(config)
	var m map[string]any
	if err := json.Unmarshal(b, &m); err != nil {
		return config, err
	}
	for k := range m {
		if v := os.Getenv(k); v != "" {
			m[k] = v
		}
	}
	b2, _ := json.Marshal(m)
	if err := json.Unmarshal(b2, &config); err != nil {
		return config, err
	}

	hello, err := base64.StdEncoding.DecodeString(config.Hello)
	if err != nil {
		return config, err
	}

	forwards := strings.Split(config.Forwards, ",")
	forwards = slices.DeleteFunc(forwards, func(s string) bool { return s == "" })
	if len(forwards) == 0 {
		return config, fmt.Errorf("at least one $FORWARD required")
	}
	config.forwards = make([]*sync.Pool, len(forwards))
	for i := range forwards {
		forward := forwards[i]
		config.forwards[i] = &sync.Pool{
			New: func() any {
				if ctx.Err() != nil {
					return nil
				}
				log.Printf("dialing %q with %q", forward, hello)
				v, err := (&net.Dialer{}).DialContext(ctx, "tcp", forward)
				if err != nil {
					log.Printf("! failed dial %q: %v", forward, err)
					return nil
				}
				if _, err := v.Write([]byte(hello)); err != nil {
					log.Printf("! failed write hello %q: %v", forward, err)
					v.Close()
					return nil
				}
				if raw, _, err := readMessage(bufio.NewReader(v)); err != nil {
					log.Printf("! failed read hello %q: %v", forward, err)
					v.Close()
					return nil
				} else {
					log.Printf("dial reply: %q", raw)
				}
				return v
			},
		}
	}

	return config, 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
}