package fwd

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"proxy-server/gateway/app"
	"proxy-server/gateway/env"
	"proxy-server/gateway/report"
	"proxy-server/utils"
	"strconv"
	"syscall"
)

type CtrlCmdType int

const (
	CtrlCmdPong CtrlCmdType = iota + 1
	CtrlCmdPing
	CtrlCmdOpen
	CtrlCmdClose
	CtrlCmdProxy
)

func (s *Service) listenCtrl() error {
	ctrlPort := env.AppCtrlPort

	// 监听端口
	ls, err := net.Listen("tcp", ":"+strconv.Itoa(int(ctrlPort)))
	if err != nil {
		return fmt.Errorf("监听控制通道失败: %w", err)
	}
	defer utils.Close(ls)

	// 处理连接
	// 异步等待连接
	var connCh = make(chan net.Conn)
	go func() {
		for {
			conn, err := ls.Accept()
			if errors.Is(err, net.ErrClosed) {
				slog.Debug("控制通道监听关闭")
				return
			}
			if err != nil {
				slog.Error("接受控制通道连接失败", "err", err)
				return
			}
			select {
			case connCh <- conn:
			case <-s.ctx.Done():
				utils.Close(conn)
				return
			}
		}
	}()

	err = nil
	for {
		select {
		case <-s.ctx.Done():
			return nil
		case conn := <-connCh:
			s.ctrlConnWg.Add(1)
			go func() {
				defer s.ctrlConnWg.Done()
				defer utils.Close(conn)
				err := s.processCtrlConn(s.ctx, conn)
				if err != nil {
					slog.Error("处理控制通道连接失败", "err", err)
				}
			}()
		}
	}
}

func (s *Service) processCtrlConn(_ctx context.Context, conn net.Conn) (err error) {
	// 通道上下文
	ctx, cancel := context.WithCancel(_ctx)

	// 结束后清理资源
	var fwdPort uint16
	defer func() {
		slog.Debug("关闭控制通道", "port", fwdPort)
		app.DelEdge(fwdPort)
	}()

	// 处理控制命令
	defer cancel()
	reader := bufio.NewReader(conn)
	for {
		// 循环等待直到服务关闭
		select {
		case <-ctx.Done():
			return nil
		default:
		}

		// 读取命令
		cmd, err := reader.ReadByte()
		if errors.Is(err, syscall.WSAECONNRESET) {
			slog.Debug("节点重置了控制通道连接(WSAECONNRESET)")
			return nil
		}
		if errors.Is(err, io.EOF) {
			slog.Debug("节点关闭了控制通道")
			return nil
		}
		if err != nil {
			return fmt.Errorf("读取节点命令失败: %w", err)
		}

		// 处理节点命令
		switch CtrlCmdType(cmd) {

		// 连接建立命令
		case CtrlCmdOpen:
			var recv = make([]byte, 4)
			_, err = io.ReadFull(reader, recv)
			if err != nil {
				return fmt.Errorf("读取节点 ID 失败: %w", err)
			}
			var client = int32(binary.BigEndian.Uint32(recv))
			fwdPort, err = s.onOpen(ctx, conn, client)
			if err != nil {
				return fmt.Errorf("处理连接建立命令失败: %w", err)
			}

		// 心跳命令
		case CtrlCmdPing:
			err = s.onPing(conn)
			if err != nil {
				return fmt.Errorf("处理心跳命令失败: %w", err)
			}

		// 连接关闭命令
		case CtrlCmdClose:
			err = s.onClose(conn)
			if err != nil {
				return fmt.Errorf("处理关闭命令失败: %w", err)
			}
			return nil

		// 忽略其他不应该由节点发起的命令
		default:
			return fmt.Errorf("无法处理控制命令: %d", cmd)
		}
	}
}

func (s *Service) onOpen(ctx context.Context, writer io.Writer, edge int32) (port uint16, err error) {
	// open 命令全局只执行一次
	_, ok := app.Edges.Load(edge)
	if ok {
		return 0, fmt.Errorf("节点 ID %d 已经连接", edge)
	}

	// 分配端口
	var minim uint16 = 20000
	var maxim uint16 = 60000
	for i := minim; i < maxim; i++ {
		var _, ok = app.Assigns.Load(i)
		if !ok {
			port = i
			app.AddEdge(edge, port)
			break
		}
	}
	if port == 0 {
		return 0, errors.New("没有可用的端口")
	}

	// 报告端口分配
	if err = report.Assigned(edge, port); err != nil {
		return 0, fmt.Errorf("报告端口分配失败: %w", err)
	}

	// 响应节点
	if err = s.sendPong(writer); err != nil {
		return 0, fmt.Errorf("响应节点失败: %w", err)
	}

	// 启动转发服务
	s.fwdLesWg.Add(1)
	go func() {
		defer s.fwdLesWg.Done()
		slog.Info("监听转发端口", "port", port, "edge", edge)
		err = s.listenUser(ctx, port, writer)
		if err != nil {
			slog.Error("监听转发端口失败", "port", port, "edge", edge, "err", err)
		}
	}()

	return port, nil
}

func (s *Service) onPing(writer io.Writer) (err error) {
	return s.sendPong(writer)
}

func (s *Service) onClose(writer io.Writer) (err error) {
	return s.sendPong(writer)
}

func (s *Service) sendPong(writer io.Writer) (err error) {
	_, err = writer.Write([]byte{byte(CtrlCmdPong)})
	if err != nil {
		return fmt.Errorf("响应节点失败: %w", err)
	}
	return nil
}

func (s *Service) sendProxy(writer io.Writer, tag [16]byte, addr string) (err error) {
	if len(addr) > 65535 {
		return fmt.Errorf("代理地址过长: %s", addr)
	}

	buf := make([]byte, 1+16+2+len(addr))
	buf[0] = byte(CtrlCmdProxy)
	copy(buf[1:], tag[:])
	binary.BigEndian.PutUint16(buf[17:], uint16(len(addr)))
	copy(buf[19:], addr)

	_, err = writer.Write(buf)
	if err != nil {
		return fmt.Errorf("发送代理命令失败: %w", err)
	}
	return nil
}