package out

import (
	"context"
	"log"
	"net"
	"net/url"
	"sync"

	"golang.org/x/net/proxy"

	"edgaru089.ink/go/regolith/internal/util"
)

// Configures upstream TCP connections.
//
// Does not support HTTP proxies as upstream. They're
// just a pain to use besides CONNECT commands.
//
// Upstreams are specified as 'direct' for direct dialing
// by the Go net package, or 'socks5://user:pass@host:port'
// for SOCKS5 upstreams.
//
// Numeric addresses are always dialed directly.
type Config struct {
	Default string              // Default URI
	Proxied map[string][]string // Proxy URI to hostname array mapping
}

// Dialer implements x/net/proxy.Dialer from the Config given to it.
type Dialer struct {
	lock sync.RWMutex

	hostnames map[string]string
	uris      map[string]interface {
		proxy.Dialer
		proxy.ContextDialer
	}

	def interface {
		proxy.Dialer
		proxy.ContextDialer
	}
}

var _ proxy.ContextDialer = &Dialer{} // *Dialer implements x/net/proxy.ContextDialer
var _ proxy.Dialer = &Dialer{}        // *Dialer implements x/net/proxy.Dialer

// New creates a new Dialer from a Config.
//
// You can also &out.Dialer{} and then call Load on it.
func New(c Config) (d *Dialer) {
	d = &Dialer{}
	d.Load(c)
	return
}

// Load loads/reloads the Dialer struct.
func (d *Dialer) Load(c Config) {
	d.lock.Lock()
	defer d.lock.Unlock()

	if len(c.Default) == 0 {
		c.Default = "direct"
	}
	log.Print("default dialer ", c.Default)

	if d.hostnames == nil {
		d.hostnames = make(map[string]string)
		d.uris = make(map[string]interface {
			proxy.Dialer
			proxy.ContextDialer
		})
	} else {
		clear(d.hostnames)
		clear(d.uris)
	}

	for uri, hosts := range c.Proxied {
		d.uris[uri] = nil
		for _, host := range hosts {
			d.hostnames[host] = uri
			log.Printf("dialer( %s ) => %s", host, uri)
		}
	}
	d.uris[c.Default] = nil

	// construct the Dialers
	for uri := range d.uris {
		if uri == "direct" {
			d.uris[uri] = proxy.Direct
		} else {
			url, err := url.Parse(uri)
			if err != nil {
				log.Print("URI parse error: ", err)
				d.uris[uri] = proxy.Direct
				continue
			}
			dx, err := proxy.FromURL(url, proxy.Direct)
			if err != nil {
				log.Print("Proxy.FromURL error: ", err)
				d.uris[uri] = proxy.Direct
				continue
			}

			var ok bool
			if d.uris[uri], ok = dx.(interface {
				proxy.ContextDialer
				proxy.Dialer
			}); !ok {
				log.Print("Proxy.FromURL unable to cast to proxy.Dialer+ContextDialer")
				d.uris[uri] = proxy.Direct
				continue
			}
		}
	}

	d.def = d.uris[c.Default]
}

func (d *Dialer) findDialer(host string) interface {
	proxy.Dialer
	proxy.ContextDialer
} {
	d.lock.RLock()
	defer d.lock.RUnlock()

	if dx, ok := d.uris[d.hostnames[host]]; ok {
		return dx
	} else {
		return d.def
	}
}

func (d *Dialer) Dial(network, address string) (net.Conn, error) {
	host, _ := util.SplitHostPort(address)
	return d.findDialer(host).Dial(network, address)
}

func (d *Dialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	host, _ := util.SplitHostPort(address)
	return d.findDialer(host).DialContext(ctx, network, address)
}