package awg

import (
	"bufio"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// DialFunc dials TCP through the AWG userspace stack.
type DialFunc func(ctx context.Context, network, address string) (net.Conn, error)

type proxyServer struct {
	mu        sync.Mutex
	socksLn   net.Listener
	httpLn    net.Listener
	dial      DialFunc
	wg        sync.WaitGroup
	closed    bool
	socksAddr string
	httpAddr  string
	conns     map[net.Conn]struct{}
}

func startProxies(socksHostPort, httpHostPort string, dial DialFunc) (*proxyServer, error) {
	s := &proxyServer{
		dial:  dial,
		conns: make(map[net.Conn]struct{}),
	}
	if socksHostPort != "" {
		ln, err := net.Listen("tcp", socksHostPort)
		if err != nil {
			return nil, fmt.Errorf("awg socks listen %s: %w", socksHostPort, err)
		}
		s.socksLn = ln
		s.socksAddr = ln.Addr().String()
		s.wg.Add(1)
		go s.serveSOCKS()
	}
	if httpHostPort != "" {
		ln, err := net.Listen("tcp", httpHostPort)
		if err != nil {
			_ = s.Close()
			return nil, fmt.Errorf("awg http listen %s: %w", httpHostPort, err)
		}
		s.httpLn = ln
		s.httpAddr = ln.Addr().String()
		s.wg.Add(1)
		go s.serveHTTP()
	}
	return s, nil
}

func (s *proxyServer) track(c net.Conn) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		_ = c.Close()
		return
	}
	s.conns[c] = struct{}{}
}

func (s *proxyServer) untrack(c net.Conn) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.conns, c)
}

func (s *proxyServer) Close() error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil
	}
	s.closed = true
	if s.socksLn != nil {
		_ = s.socksLn.Close()
	}
	if s.httpLn != nil {
		_ = s.httpLn.Close()
	}
	// Force-close active relays so io.Copy / dial wake up immediately.
	for c := range s.conns {
		_ = c.Close()
	}
	s.conns = nil
	s.mu.Unlock()

	done := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(1500 * time.Millisecond):
		// Don't hang UI / Disconnect on stuck netstack copies.
	}
	return nil
}

func (s *proxyServer) serveSOCKS() {
	defer s.wg.Done()
	for {
		c, err := s.socksLn.Accept()
		if err != nil {
			return
		}
		s.wg.Add(1)
		go func(conn net.Conn) {
			defer s.wg.Done()
			s.track(conn)
			defer s.untrack(conn)
			_ = handleSOCKS(conn, s.dial)
		}(c)
	}
}

func (s *proxyServer) serveHTTP() {
	defer s.wg.Done()
	for {
		c, err := s.httpLn.Accept()
		if err != nil {
			return
		}
		s.wg.Add(1)
		go func(conn net.Conn) {
			defer s.wg.Done()
			s.track(conn)
			defer s.untrack(conn)
			_ = handleHTTPProxy(conn, s.dial)
		}(c)
	}
}

func handleSOCKS(client net.Conn, dial DialFunc) error {
	defer client.Close()
	_ = client.SetDeadline(time.Now().Add(30 * time.Second))
	buf := make([]byte, 258)
	if _, err := io.ReadFull(client, buf[:2]); err != nil {
		return err
	}
	if buf[0] != 0x05 {
		return fmt.Errorf("socks: not v5")
	}
	nmethods := int(buf[1])
	if _, err := io.ReadFull(client, buf[:nmethods]); err != nil {
		return err
	}
	if _, err := client.Write([]byte{0x05, 0x00}); err != nil {
		return err
	}
	if _, err := io.ReadFull(client, buf[:4]); err != nil {
		return err
	}
	if buf[0] != 0x05 || buf[1] != 0x01 {
		_, _ = client.Write([]byte{0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return fmt.Errorf("socks: only CONNECT supported")
	}
	var host string
	switch buf[3] {
	case 0x01: // IPv4
		if _, err := io.ReadFull(client, buf[:4]); err != nil {
			return err
		}
		host = net.IP(buf[:4]).String()
	case 0x03: // domain
		if _, err := io.ReadFull(client, buf[:1]); err != nil {
			return err
		}
		l := int(buf[0])
		if _, err := io.ReadFull(client, buf[:l]); err != nil {
			return err
		}
		host = string(buf[:l])
	case 0x04: // IPv6
		if _, err := io.ReadFull(client, buf[:16]); err != nil {
			return err
		}
		host = net.IP(buf[:16]).String()
	default:
		_, _ = client.Write([]byte{0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return fmt.Errorf("socks: bad atyp")
	}
	if _, err := io.ReadFull(client, buf[:2]); err != nil {
		return err
	}
	port := binary.BigEndian.Uint16(buf[:2])
	_ = client.SetDeadline(time.Time{})

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	remote, err := dial(ctx, "tcp", net.JoinHostPort(host, fmt.Sprintf("%d", port)))
	if err != nil {
		_, _ = client.Write([]byte{0x05, 0x05, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return err
	}
	defer remote.Close()
	if _, err := client.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}); err != nil {
		return err
	}
	return relay(client, remote)
}

func handleHTTPProxy(client net.Conn, dial DialFunc) error {
	defer client.Close()
	_ = client.SetDeadline(time.Now().Add(30 * time.Second))
	br := bufio.NewReader(client)
	req, err := http.ReadRequest(br)
	if err != nil {
		return err
	}
	_ = client.SetDeadline(time.Time{})

	if strings.EqualFold(req.Method, http.MethodConnect) {
		host := req.Host
		if !strings.Contains(host, ":") {
			host += ":443"
		}
		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
		defer cancel()
		remote, err := dial(ctx, "tcp", host)
		if err != nil {
			_, _ = client.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
			return err
		}
		defer remote.Close()
		if _, err := client.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil {
			return err
		}
		return relay(client, remote)
	}

	// Absolute-form HTTP proxy request.
	u := req.URL
	host := u.Host
	if host == "" {
		host = req.Host
	}
	if host == "" {
		_, _ = client.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n"))
		return fmt.Errorf("http proxy: no host")
	}
	if !strings.Contains(host, ":") {
		host += ":80"
	}
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	remote, err := dial(ctx, "tcp", host)
	if err != nil {
		_, _ = client.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
		return err
	}
	defer remote.Close()
	req.RequestURI = ""
	if err := req.Write(remote); err != nil {
		return err
	}
	return relay(client, remote)
}

func relay(a, b net.Conn) error {
	errc := make(chan error, 2)
	copyOne := func(dst, src net.Conn) {
		bufp := copyBufPool.Get().(*[]byte)
		defer copyBufPool.Put(bufp)
		_, err := io.CopyBuffer(dst, src, *bufp)
		errc <- err
	}
	go copyOne(a, b)
	go copyOne(b, a)
	err := <-errc
	_ = a.Close()
	_ = b.Close()
	<-errc
	return err
}

var copyBufPool = sync.Pool{
	New: func() any {
		b := make([]byte, 32*1024)
		return &b
	},
}