package socks

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"net"
	"proxy-server/pkg/utils"
	"strconv"

	"github.com/pkg/errors"
)

const (
	ConnectCommand   = byte(1)
	BindCommand      = byte(2)
	AssociateCommand = byte(3)
	ipv4Address      = byte(1)
	fqdnAddress      = byte(3)
	ipv6Address      = byte(4)
)

const (
	successReply byte = iota
	serverFailure
	ruleFailure
	networkUnreachable
	hostUnreachable
	connectionRefused
	ttlExpired
	commandNotSupported
	addrTypeNotSupported
)

var (
	unrecognizedAddrType = fmt.Errorf("unrecognized address type")
)

// AddressRewriter is used to rewrite a destination transparently
type AddressRewriter interface {
	Rewrite(ctx context.Context, request *Request) (context.Context, *AddrSpec)
}

// AddrSpec 地址
type AddrSpec struct {
	FQDN string
	IP   net.IP
	Port int
}

func (a AddrSpec) String() string {
	if a.FQDN != "" {
		return fmt.Sprintf("%s (%s):%d", a.FQDN, a.IP, a.Port)
	}
	return fmt.Sprintf("%s:%d", a.IP, a.Port)
}

// Address returns a string suitable to dial; prefer returning IP-based
// address, fallback to FQDN
func (a AddrSpec) Address() string {
	if 0 != len(a.IP) {
		return net.JoinHostPort(a.IP.String(), strconv.Itoa(a.Port))
	}
	return net.JoinHostPort(a.FQDN, strconv.Itoa(a.Port))
}

func (s *Server) request(reader io.Reader, writer io.Writer) (*Request, error) {

	// 检查版本
	err := checkVersion(reader)
	if err != nil {
		return nil, err
	}

	// 检查连接命令
	command, err := utils.ReadByte(reader)
	if err != nil {
		return nil, err
	}

	slog.Debug("客户端使用的连接指令：%v", command)
	if command != ConnectCommand && command != BindCommand && command != AssociateCommand {
		err = sendReply(writer, commandNotSupported, nil)
		if err != nil {
			return nil, err
		}
		return nil, errors.New("不支持该连接指令")
	}

	// 跳过保留字段 rsv
	_, err = utils.ReadByte(reader)
	if err != nil {
		return nil, err
	}

	// 获取目标地址
	dest, err := s.parseTarget(reader, writer)
	if err != nil {
		return nil, err
	}

	request := &Request{
		Version:  Version,
		Command:  command,
		DestAddr: dest,
		bufConn:  reader,
	}

	return request, nil
}

func (s *Server) parseTarget(reader io.Reader, writer io.Writer) (*AddrSpec, error) {
	dest := &AddrSpec{}

	aTypeBuf := make([]byte, 1)
	_, err := reader.Read(aTypeBuf)
	if err != nil {
		return nil, err
	}

	switch aTypeBuf[0] {

	case ipv4Address:
		addr := make([]byte, 4)
		_, err := io.ReadFull(reader, addr)
		if err != nil {
			return nil, err
		}
		dest.IP = addr

	case ipv6Address:
		addr := make([]byte, 16)
		_, err := io.ReadFull(reader, addr)
		if err != nil {
			return nil, err
		}
		dest.IP = addr

	case fqdnAddress:
		aLenBuf := make([]byte, 1)
		_, err := reader.Read(aLenBuf)
		if err != nil {
			return nil, err
		}

		fqdnBuff := make([]byte, int(aLenBuf[0]))
		_, err = io.ReadFull(reader, fqdnBuff)
		if err != nil {
			return nil, err
		}
		dest.FQDN = string(fqdnBuff)

		// 域名解析
		addr, err := s.config.Resolver.Resolve(dest.FQDN)
		if err != nil {
			err := sendReply(writer, hostUnreachable, nil)
			if err != nil {
				return nil, fmt.Errorf("failed to send reply: %v", err)
			}
			return nil, fmt.Errorf("failed to resolve destination '%v': %v", dest.FQDN, err)
		}
		dest.IP = addr

	default:
		err := sendReply(writer, addrTypeNotSupported, nil)
		if err != nil {
			return nil, err
		}
		return nil, unrecognizedAddrType
	}

	portBuf := make([]byte, 2)
	_, err = io.ReadFull(reader, portBuf)
	if err != nil {
		return nil, err
	}
	dest.Port = (int(portBuf[0]) << 8) | int(portBuf[1])

	return dest, nil
}

// A Request represents request received by a server
type Request struct {
	// Protocol version
	Version uint8
	// Requested command
	Command uint8
	// Authentication provided during negotiation
	Authentication *Authentication
	// AddrSpec of the  network that sent the request
	RemoteAddr *AddrSpec
	// AddrSpec of the desired destination
	DestAddr *AddrSpec
	// AddrSpec of the actual destination (might be affected by rewrite)
	realDestAddr *AddrSpec
	bufConn      io.Reader
}

func (s *Server) handle(req *Request, conn net.Conn) error {
	ctx := context.Background()

	// 目标地址重写
	req.realDestAddr = req.DestAddr
	if s.config.Rewriter != nil {
		ctx, req.realDestAddr = s.config.Rewriter.Rewrite(ctx, req)
	}

	// 根据协商方法建立连接
	switch req.Command {
	case ConnectCommand:
		return s.handleConnect(ctx, conn, req)
	case BindCommand:
		return s.handleBind(ctx, conn, req)
	case AssociateCommand:
		return s.handleAssociate(ctx, conn, req)
	default:
		return fmt.Errorf("unsupported command: %v", req.Command)
	}
}

func (s *Server) handleConnect(ctx context.Context, conn net.Conn, req *Request) error {
	// 检查规则集约束
	s.config.Logger.Printf("检查约束规则\n")
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("failed to send reply: %v", err)
		}
		return fmt.Errorf("request to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	slog.Info("需要向 " + req.DestAddr.Address() + " 建立连接")
	select {
	case <-s.ctx.Done():
		if conn != nil {
			utils.Close(conn)
		}
	case s.Conn <- ProxyConn{
		req.Authentication.Payload.ID,
		conn,
		req.realDestAddr.Address(),
	}:
	}
	return nil
}

func (s *Server) handleBind(ctx context.Context, conn net.Conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("failed to send reply: %v", err)
		}
		return fmt.Errorf("bind to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// TODO: Support bind
	if err := sendReply(conn, commandNotSupported, nil); err != nil {
		return fmt.Errorf("failed to send reply: %v", err)
	}
	return nil
}

func (s *Server) handleAssociate(ctx context.Context, conn net.Conn, req *Request) error {
	// Check if this is allowed
	if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
		if err := sendReply(conn, ruleFailure, nil); err != nil {
			return fmt.Errorf("failed to send reply: %v", err)
		}
		return fmt.Errorf("associate to %v blocked by rules", req.DestAddr)
	} else {
		ctx = ctx_
	}

	// TODO: Support associate
	if err := sendReply(conn, commandNotSupported, nil); err != nil {
		return fmt.Errorf("failed to send reply: %v", err)
	}
	return nil
}

func sendReply(w io.Writer, resp uint8, addr *AddrSpec) error {
	var addrType uint8
	var addrBody []byte
	var addrPort uint16
	switch {
	case addr == nil:
		addrType = ipv4Address
		addrBody = []byte{0, 0, 0, 0}
		addrPort = 0

	case addr.FQDN != "":
		addrType = fqdnAddress
		addrBody = append([]byte{byte(len(addr.FQDN))}, addr.FQDN...)
		addrPort = uint16(addr.Port)

	case addr.IP.To4() != nil:
		addrType = ipv4Address
		addrBody = addr.IP.To4()
		addrPort = uint16(addr.Port)

	case addr.IP.To16() != nil:
		addrType = ipv6Address
		addrBody = addr.IP.To16()
		addrPort = uint16(addr.Port)

	default:
		return fmt.Errorf("failed to format address: %v", addr)
	}

	msg := make([]byte, 6+len(addrBody))
	msg[0] = Version
	msg[1] = resp
	msg[2] = 0 // Reserved
	msg[3] = addrType
	copy(msg[4:], addrBody)
	msg[4+len(addrBody)] = byte(addrPort >> 8)
	msg[4+len(addrBody)+1] = byte(addrPort & 0xff)

	_, err := w.Write(msg)
	return err
}

func SendSuccess(user net.Conn, target net.Conn) error {
	local := target.LocalAddr().(*net.TCPAddr)
	bind := AddrSpec{IP: local.IP, Port: local.Port}
	err := sendReply(user, successReply, &bind)
	if err != nil {
		return err
	}
	return nil
}

type ProxyConn struct {
	Uid uint
	// 用户连入的连接
	Conn net.Conn
	// 用户目标地址
	Dest string
}

func (d ProxyConn) Tag() string {
	local := d.Conn.LocalAddr()
	remote := d.Conn.RemoteAddr()
	return fmt.Sprintf("%s-%s", remote, local)
}

func (d ProxyConn) Close() error {
	return d.Conn.Close()
}