package handler

import (
	"fmt"
	"io"
	"net"
	"sync"

	"git.kingecg.top/kingecg/gologger"
)

// TcpProxyHandler TCP代理处理器
type TcpProxyHandler struct {
	Name     string
	Listen   string
	Target   string
	listener net.Listener
	wg       sync.WaitGroup
	stopCh   chan struct{}
}

func NewTcpProxyHandler(name, listen, target string) *TcpProxyHandler {
	return &TcpProxyHandler{
		Name:   name,
		Listen: listen,
		Target: target,
		stopCh: make(chan struct{}),
	}
}

// Start 开始监听并处理TCP连接
func (t *TcpProxyHandler) Start() error {
	l := gologger.GetLogger("TcpProxy")
	l.Info(fmt.Sprintf("TCP proxy [%s] starting, listen: %s, target: %s", t.Name, t.Listen, t.Target))

	var err error
	t.listener, err = net.Listen("tcp", t.Listen)
	if err != nil {
		l.Error(fmt.Sprintf("TCP proxy [%s] failed to listen on %s: %v", t.Name, t.Listen, err))
		return err
	}

	go t.acceptLoop()
	return nil
}

// acceptLoop 接受连接并处理
func (t *TcpProxyHandler) acceptLoop() {
	l := gologger.GetLogger("TcpProxy")
	for {
		conn, err := t.listener.Accept()
		if err != nil {
			select {
			case <-t.stopCh:
				l.Info(fmt.Sprintf("TCP proxy [%s] stopped", t.Name))
				return
			default:
				l.Error(fmt.Sprintf("TCP proxy [%s] accept error: %v", t.Name, err))
				continue
			}
		}
		t.wg.Add(1)
		go t.handleConn(conn)
	}
}

// handleConn 处理单个TCP连接
func (t *TcpProxyHandler) handleConn(clientConn net.Conn) {
	defer t.wg.Done()
	l := gologger.GetLogger("TcpProxy")

	l.Debug(fmt.Sprintf("TCP proxy [%s] new connection from %s", t.Name, clientConn.RemoteAddr()))

	// 连接到目标服务器
	targetConn, err := net.DialTimeout("tcp", t.Target, 10*1000000000) // 10秒超时
	if err != nil {
		l.Error(fmt.Sprintf("TCP proxy [%s] failed to connect to target %s: %v", t.Name, t.Target, err))
		clientConn.Close()
		return
	}

	// 使用管道模式转发数据
	go t.pipe(clientConn, targetConn, "client->target")
	go t.pipe(targetConn, clientConn, "target->client")

	l.Debug(fmt.Sprintf("TCP proxy [%s] connection established", t.Name))
}

// pipe 在两个连接之间转发数据
func (t *TcpProxyHandler) pipe(dst, src net.Conn, direction string) {
	defer dst.Close()
	l := gologger.GetLogger("TcpProxy")
	buf := make([]byte, 8192)
	for {
		n, err := src.Read(buf)
		if err != nil {
			if err != io.EOF {
				l.Debug(fmt.Sprintf("TCP proxy [%s] %s read error: %v", t.Name, direction, err))
			}
			return
		}
		_, err = dst.Write(buf[:n])
		if err != nil {
			l.Debug(fmt.Sprintf("TCP proxy [%s] %s write error: %v", t.Name, direction, err))
			return
		}
	}
}

// Stop 停止TCP代理服务器
func (t *TcpProxyHandler) Stop() error {
	l := gologger.GetLogger("TcpProxy")
	l.Info(fmt.Sprintf("TCP proxy [%s] stopping...", t.Name))
	close(t.stopCh)
	if t.listener != nil {
		t.listener.Close()
	}
	t.wg.Wait()
	l.Info(fmt.Sprintf("TCP proxy [%s] stopped", t.Name))
	return nil
}