package main

import (
	"flag"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"os"
	"time"

	"github.com/darkit/goproxy"
	"github.com/darkit/goproxy/pkg/dns"
)

func main() {
	// 解析命令行参数
	var (
		listenAddr     = flag.String("listen", ":8080", "代理服务器监听地址")
		targetAddr     = flag.String("target", "http://example.com", "目标服务器地址")
		useDNSResolver = flag.Bool("custom-dns", false, "是否使用自定义DNS解析")
		dnsHost        = flag.String("dns-host", "example.com", "要解析的主机名")
		dnsIP          = flag.String("dns-ip", "127.0.0.1", "解析的IP地址")
		verbosity      = flag.Int("v", 0, "日志详细级别 (0-3)")
	)
	flag.Parse()

	// 配置日志
	logLevel := slog.LevelInfo
	switch *verbosity {
	case 1:
		logLevel = slog.LevelDebug
	case 2:
		logLevel = slog.LevelInfo
	case 3:
		logLevel = slog.LevelDebug - 3 // 更详细的调试级别
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
		Level: logLevel,
	}))
	slog.SetDefault(logger)

	// 创建反向代理实例
	var proxy goproxy.UnifiedProxy
	var err error

	if *useDNSResolver {
		// 使用自定义DNS解析器
		slog.Info("使用自定义DNS解析器")

		// 创建自定义DNS解析器
		resolver := dns.NewResolver(
			dns.WithFallback(true),     // 如果找不到自定义规则，回退到系统DNS
			dns.WithTTL(5*time.Minute), // 设置缓存TTL
		)

		// 添加自定义DNS解析规则
		if err := resolver.Add(*dnsHost, *dnsIP); err != nil {
			slog.Error("添加DNS解析规则失败", "error", err)
			return
		}

		slog.Info("添加DNS解析规则", "host", *dnsHost, "ip", *dnsIP)

		// 创建带自定义DNS解析器的反向代理
		proxy, err = goproxy.NewReverseProxy(*listenAddr, *targetAddr, goproxy.WithUnifiedDNSResolver(resolver))
	} else {
		// 使用系统DNS解析
		slog.Info("使用系统DNS解析")
		proxy, err = goproxy.NewReverseProxy(*listenAddr, *targetAddr)
	}

	if err != nil {
		slog.Error("创建代理服务器失败", "error", err)
		return
	}

	// 使用系统DNS解析一次目标地址，打印结果供比较
	targetHost := *targetAddr
	// 尝试解析URL
	if u, err := url.Parse(targetHost); err == nil && u.Host != "" {
		targetHost = u.Host
	}
	// 如果还包含端口，移除端口
	if host, _, err := net.SplitHostPort(targetHost); err == nil {
		targetHost = host
	}
	ips, err := net.LookupIP(targetHost)
	if err != nil {
		slog.Error("系统DNS解析失败", "host", targetHost, "error", err)
	} else {
		slog.Info("系统DNS解析结果", "host", targetHost, "ips", ips)
	}

	// 创建HTTP服务器
	server := &http.Server{
		Addr:    *listenAddr,
		Handler: proxy,
	}

	// 启动服务器
	slog.Info("代理服务器已启动",
		"listen", *listenAddr,
		"target", *targetAddr,
		"custom_dns", *useDNSResolver)

	// 监听和服务连接
	if err := server.ListenAndServe(); err != nil {
		slog.Error("服务器运行失败", "error", err)
	}
}