package main

import (
	"context"
	"flag"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"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", "目标服务器地址（反向代理模式下使用）")
		proxyMode   = flag.String("mode", "reverse", "代理模式: forward, reverse, transparent")
		enableHTTPS = flag.Bool("https", false, "是否启用HTTPS拦截")
		verbosity   = flag.Int("v", 0, "日志详细级别 (0-3)")
		host        = flag.String("dns-host", "example.com", "要解析的主机名")
		ip          = flag.String("dns-ip", "127.0.0.1", "解析的IP地址")
	)
	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)

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

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

	// 添加通配符DNS解析规则示例
	if err := resolver.AddWildcard("*.example.org", "192.168.1.2"); err != nil {
		slog.Error("添加通配符DNS解析规则失败", "error", err)
		return
	}

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

	// 创建代理选项
	var (
		proxy goproxy.UnifiedProxy
		err   error
	)

	// 根据代理模式创建不同类型的代理
	switch *proxyMode {
	case "forward":
		proxy, err = goproxy.NewForwardProxy(*listenAddr, goproxy.WithUnifiedDNSResolver(resolver))
	case "reverse":
		proxy, err = goproxy.NewReverseProxy(*listenAddr, *targetAddr, goproxy.WithUnifiedDNSResolver(resolver))
	case "transparent":
		proxy, err = goproxy.NewTransparentProxy(*listenAddr, goproxy.WithUnifiedDNSResolver(resolver))
	default:
		slog.Error("不支持的代理模式", "mode", *proxyMode)
		return
	}

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

	// 设置信号处理
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	// 打印启动信息
	slog.Info("代理服务器已启动",
		"mode", *proxyMode,
		"listen", *listenAddr,
		"target", *targetAddr,
		"https", *enableHTTPS)

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

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

	// 等待中断信号
	<-sigCh
	slog.Info("接收到信号，正在关闭...")

	// 创建关闭上下文
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// 优雅关闭HTTP服务器
	if err := server.Shutdown(ctx); err != nil {
		slog.Error("服务器关闭失败", "error", err)
	}

	// 关闭代理
	if err := proxy.Close(); err != nil {
		slog.Error("关闭代理时出错", "error", err)
	}
}