package socks5proxy

import (
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
	"github.com/shadow1ng/fscan/plugins/local"
)

// Socks5ProxyPlugin SOCKS5代理插件
type Socks5ProxyPlugin struct {
	*local.BaseLocalPlugin
	connector *Socks5ProxyConnector
	port      int
	listener  net.Listener
}

// Socks5ProxyConnector SOCKS5代理连接器
type Socks5ProxyConnector struct {
	*local.BaseLocalConnector
	port int
}

// Socks5ProxyConnection SOCKS5代理连接对象
type Socks5ProxyConnection struct {
	*local.LocalConnection
	Port int
}

// NewSocks5ProxyPlugin 创建SOCKS5代理插件
func NewSocks5ProxyPlugin() *Socks5ProxyPlugin {
	// 从全局参数获取SOCKS5端口
	port := common.Socks5ProxyPort
	if port <= 0 {
		port = 1080 // 默认端口
	}

	metadata := &base.PluginMetadata{
		Name:        "socks5proxy",
		Version:     "1.0.0",
		Author:      "fscan-team",
		Description: "本地SOCKS5代理服务器插件，支持HTTP/HTTPS代理",
		Category:    "local",
		Tags:        []string{"local", "proxy", "socks5", "network"},
		Protocols:   []string{"local"},
	}

	connector := NewSocks5ProxyConnector(port)
	plugin := &Socks5ProxyPlugin{
		BaseLocalPlugin: local.NewBaseLocalPlugin(metadata, connector),
		connector:       connector,
		port:            port,
	}

	// 设置支持的平台
	plugin.SetPlatformSupport([]string{"windows", "linux", "darwin"})
	// 不需要特殊权限
	plugin.SetRequiresPrivileges(false)

	return plugin
}

// NewSocks5ProxyConnector 创建SOCKS5代理连接器
func NewSocks5ProxyConnector(port int) *Socks5ProxyConnector {
	baseConnector, _ := local.NewBaseLocalConnector()
	
	return &Socks5ProxyConnector{
		BaseLocalConnector: baseConnector,
		port:               port,
	}
}

// Connect 建立SOCKS5代理连接
func (c *Socks5ProxyConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	// 先建立基础本地连接
	localConn, err := c.BaseLocalConnector.Connect(ctx, info)
	if err != nil {
		return nil, err
	}
	
	baseConn := localConn.(*local.LocalConnection)
	
	// 检查端口是否可用
	testListener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", c.port))
	if err != nil {
		return nil, fmt.Errorf("端口 %d 不可用: %v", c.port, err)
	}
	testListener.Close() // 立即关闭测试监听器
	
	socks5Conn := &Socks5ProxyConnection{
		LocalConnection: baseConn,
		Port:            c.port,
	}
	
	return socks5Conn, nil
}

// Close 关闭SOCKS5代理连接
func (c *Socks5ProxyConnector) Close(conn interface{}) error {
	return c.BaseLocalConnector.Close(conn)
}

// Scan 重写扫描方法以确保调用正确的ScanLocal实现
func (p *Socks5ProxyPlugin) Scan(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	return p.ScanLocal(ctx, info)
}

// ScanLocal 执行SOCKS5代理扫描 - 启动代理服务器
func (p *Socks5ProxyPlugin) ScanLocal(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	common.LogBase("启动SOCKS5代理服务器...")
	
	// 建立连接
	conn, err := p.connector.Connect(ctx, info)
	if err != nil {
		return &base.ScanResult{
			Success: false,
			Error:   fmt.Errorf("连接失败: %v", err),
		}, nil
	}
	defer p.connector.Close(conn)
	
	socks5Conn := conn.(*Socks5ProxyConnection)
	
	// 启动SOCKS5代理服务器
	common.LogBase(fmt.Sprintf("在端口 %d 上启动SOCKS5代理", socks5Conn.Port))
	
	// 直接在当前goroutine中运行，这样可以确保在设置Socks5ProxyActive后立即被主程序检测到
	err = p.startSocks5Server(ctx, socks5Conn.Port)
	if err != nil {
		common.LogError(fmt.Sprintf("SOCKS5代理服务器错误: %v", err))
		return &base.ScanResult{
			Success: false,
			Error:   err,
		}, nil
	}
	
	result := &base.ScanResult{
		Success: true,
		Service: "SOCKS5Proxy",
		Banner:  fmt.Sprintf("SOCKS5代理已完成 - 端口: %d 平台: %s", socks5Conn.Port, runtime.GOOS),
		Extra: map[string]interface{}{
			"port":     socks5Conn.Port,
			"platform": runtime.GOOS,
			"protocol": "socks5",
			"status":   "completed",
		},
	}
	
	return result, nil
}

// startSocks5Server 启动SOCKS5代理服务器 - 核心实现
func (p *Socks5ProxyPlugin) startSocks5Server(ctx context.Context, port int) error {
	// 监听指定端口
	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port))
	if err != nil {
		return fmt.Errorf("监听端口失败: %v", err)
	}
	defer listener.Close()
	
	p.listener = listener
	common.LogSuccess(fmt.Sprintf("SOCKS5代理服务器已在 127.0.0.1:%d 上启动", port))
	
	// 设置SOCKS5代理为活跃状态，告诉主程序保持运行
	common.Socks5ProxyActive = true
	defer func() {
		// 确保退出时清除活跃状态
		common.Socks5ProxyActive = false
	}()
	
	// 主循环处理连接
	for {
		select {
		case <-ctx.Done():
			common.LogBase("SOCKS5代理服务器被上下文取消")
			return ctx.Err()
		default:
		}
		
		// 设置监听器超时，以便能响应上下文取消
		if tcpListener, ok := listener.(*net.TCPListener); ok {
			tcpListener.SetDeadline(time.Now().Add(1 * time.Second))
		}
		
		conn, err := listener.Accept()
		if err != nil {
			// 检查是否是超时错误
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue // 超时继续循环
			}
			common.LogError(fmt.Sprintf("接受连接失败: %v", err))
			continue
		}
		
		// 并发处理客户端连接
		go p.handleClient(conn)
	}
}

// handleClient 处理客户端连接
func (p *Socks5ProxyPlugin) handleClient(clientConn net.Conn) {
	defer clientConn.Close()
	
	// SOCKS5握手阶段
	if err := p.handleSocks5Handshake(clientConn); err != nil {
		common.LogError(fmt.Sprintf("SOCKS5握手失败: %v", err))
		return
	}
	
	// SOCKS5请求阶段
	targetConn, err := p.handleSocks5Request(clientConn)
	if err != nil {
		common.LogError(fmt.Sprintf("SOCKS5请求处理失败: %v", err))
		return
	}
	defer targetConn.Close()
	
	common.LogSuccess("建立SOCKS5代理连接")
	
	// 双向数据转发
	p.relayData(clientConn, targetConn)
}

// handleSocks5Handshake 处理SOCKS5握手
func (p *Socks5ProxyPlugin) handleSocks5Handshake(conn net.Conn) error {
	// 读取客户端握手请求
	buffer := make([]byte, 256)
	n, err := conn.Read(buffer)
	if err != nil {
		return fmt.Errorf("读取握手请求失败: %v", err)
	}
	
	if n < 3 || buffer[0] != 0x05 { // SOCKS版本必须是5
		return fmt.Errorf("不支持的SOCKS版本")
	}
	
	// 发送握手响应（无认证）
	response := []byte{0x05, 0x00} // 版本5，无认证
	_, err = conn.Write(response)
	if err != nil {
		return fmt.Errorf("发送握手响应失败: %v", err)
	}
	
	return nil
}

// handleSocks5Request 处理SOCKS5连接请求
func (p *Socks5ProxyPlugin) handleSocks5Request(clientConn net.Conn) (net.Conn, error) {
	// 读取连接请求
	buffer := make([]byte, 256)
	n, err := clientConn.Read(buffer)
	if err != nil {
		return nil, fmt.Errorf("读取连接请求失败: %v", err)
	}
	
	if n < 7 || buffer[0] != 0x05 {
		return nil, fmt.Errorf("无效的SOCKS5请求")
	}
	
	cmd := buffer[1]
	if cmd != 0x01 { // 只支持CONNECT命令
		// 发送不支持的命令响应
		response := []byte{0x05, 0x07, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
		clientConn.Write(response)
		return nil, fmt.Errorf("不支持的命令: %d", cmd)
	}
	
	// 解析目标地址
	addrType := buffer[3]
	var targetHost string
	var targetPort int
	
	switch addrType {
	case 0x01: // IPv4
		if n < 10 {
			return nil, fmt.Errorf("IPv4地址格式错误")
		}
		targetHost = fmt.Sprintf("%d.%d.%d.%d", buffer[4], buffer[5], buffer[6], buffer[7])
		targetPort = int(buffer[8])<<8 + int(buffer[9])
	case 0x03: // 域名
		if n < 5 {
			return nil, fmt.Errorf("域名格式错误")
		}
		domainLen := int(buffer[4])
		if n < 5+domainLen+2 {
			return nil, fmt.Errorf("域名长度错误")
		}
		targetHost = string(buffer[5 : 5+domainLen])
		targetPort = int(buffer[5+domainLen])<<8 + int(buffer[5+domainLen+1])
	case 0x04: // IPv6
		if n < 22 {
			return nil, fmt.Errorf("IPv6地址格式错误")
		}
		// IPv6地址解析（简化实现）
		targetHost = net.IP(buffer[4:20]).String()
		targetPort = int(buffer[20])<<8 + int(buffer[21])
	default:
		// 发送不支持的地址类型响应
		response := []byte{0x05, 0x08, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
		clientConn.Write(response)
		return nil, fmt.Errorf("不支持的地址类型: %d", addrType)
	}
	
	// 连接目标服务器
	targetAddr := fmt.Sprintf("%s:%d", targetHost, targetPort)
	targetConn, err := net.DialTimeout("tcp", targetAddr, 10*time.Second)
	if err != nil {
		// 发送连接失败响应
		response := []byte{0x05, 0x05, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
		clientConn.Write(response)
		return nil, fmt.Errorf("连接目标服务器失败: %v", err)
	}
	
	// 发送成功响应
	response := make([]byte, 10)
	response[0] = 0x05 // SOCKS版本
	response[1] = 0x00 // 成功
	response[2] = 0x00 // 保留
	response[3] = 0x01 // IPv4地址类型
	// 绑定地址和端口（使用127.0.0.1:port）
	copy(response[4:8], []byte{127, 0, 0, 1})
	response[8] = byte(p.port >> 8)
	response[9] = byte(p.port & 0xff)
	
	_, err = clientConn.Write(response)
	if err != nil {
		targetConn.Close()
		return nil, fmt.Errorf("发送成功响应失败: %v", err)
	}
	
	common.LogDebug(fmt.Sprintf("建立代理连接: %s", targetAddr))
	return targetConn, nil
}

// relayData 双向数据转发
func (p *Socks5ProxyPlugin) relayData(clientConn, targetConn net.Conn) {
	done := make(chan struct{}, 2)
	
	// 客户端到目标服务器
	go func() {
		defer func() { done <- struct{}{} }()
		io.Copy(targetConn, clientConn)
		targetConn.Close()
	}()
	
	// 目标服务器到客户端
	go func() {
		defer func() { done <- struct{}{} }()
		io.Copy(clientConn, targetConn)
		clientConn.Close()
	}()
	
	// 等待其中一个方向完成
	<-done
}

// GetLocalData 获取SOCKS5代理本地数据
func (p *Socks5ProxyPlugin) GetLocalData(ctx context.Context) (map[string]interface{}, error) {
	data := make(map[string]interface{})
	
	// 获取系统信息
	data["plugin_type"] = "socks5proxy"
	data["platform"] = runtime.GOOS
	data["arch"] = runtime.GOARCH
	data["port"] = p.port
	data["protocol"] = "socks5"
	
	if homeDir, err := os.UserHomeDir(); err == nil {
		data["home_dir"] = homeDir
	}
	
	if workDir, err := os.Getwd(); err == nil {
		data["work_dir"] = workDir
	}
	
	return data, nil
}

// ExtractData 提取数据（SOCKS5代理主要是服务功能）
func (p *Socks5ProxyPlugin) ExtractData(ctx context.Context, info *common.HostInfo, data map[string]interface{}) (*base.ExploitResult, error) {
	return &base.ExploitResult{
		Success: true,
		Output:  fmt.Sprintf("SOCKS5代理服务器运行完成，端口: %d", p.port),
		Data:    data,
		Extra: map[string]interface{}{
			"port":     p.port,
			"protocol": "socks5",
			"status":   "completed",
		},
	}, nil
}

// GetInfo 获取插件信息
func (p *Socks5ProxyPlugin) GetInfo() string {
	var info strings.Builder
	
	info.WriteString(fmt.Sprintf("SOCKS5代理插件 - 端口: %d\n", p.port))
	info.WriteString(fmt.Sprintf("支持平台: %s\n", strings.Join(p.GetPlatformSupport(), ", ")))
	info.WriteString("协议: SOCKS5，支持HTTP/HTTPS代理\n")
	info.WriteString("实现方式: 纯Go原生，无外部依赖\n")
	
	return info.String()
}

// RegisterSocks5ProxyPlugin 注册SOCKS5代理插件
func RegisterSocks5ProxyPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "socks5proxy",
			Version:     "1.0.0", 
			Author:      "fscan-team",
			Description: "本地SOCKS5代理服务器插件，支持HTTP/HTTPS代理",
			Category:    "local",
			Tags:        []string{"socks5proxy", "local", "proxy", "network"},
			Protocols:   []string{"local"},
		},
		func() base.Plugin {
			return NewSocks5ProxyPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("socks5proxy", factory)
}

// init 插件注册函数
func init() {
	RegisterSocks5ProxyPlugin()
}