package main

import (
	"context"
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

type Config struct {
	LuckyAPI       string `json:"lucky_api"`
	LuckyAPITarget string `json:"lucky_api_target"`
	ConfigFile     string `json:"-"`
	Rules          []Rule `json:"rules"`
	Proxy          string `json:"proxy"`
	Cert           string `json:"cert"`
	Key            string `json:"key"`
	SmartDNSIP     string `json:"smartdns_ip"`
	SmartDSConf   string `json:"smartdns_conf"`
	SmartDSReload string `json:"smartdns_reload"`
	Interval       string `json:"interval"`
}

type Rule struct {
	Host   string `json:"host"`
	Target string `json:"target"`
}

var (
	cfg            Config
	allowedDomains map[string]string // 域名 → target
	apiDomains    map[string]string // Lucky API 域名 → target
	fileDomains    map[string]string // config.json rules 域名 → target
	cert           tls.Certificate
	configModTime  time.Time
)

func main() {
	// 调大文件描述符限制
	var rlimit syscall.Rlimit
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rlimit); err == nil {
		rlimit.Cur = rlimit.Max
		syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rlimit)
	}

	configFlag := flag.String("config", "", "JSON 配置文件路径")
	flag.Parse()

	if *configFlag != "" {
		cfg.ConfigFile = *configFlag
		if err := loadConfig(*configFlag); err != nil {
			log.Fatalf("[Config] 加载失败: %v", err)
		}
	}

	// 加载证书
	if err := loadCert(); err != nil {
		log.Printf("[Cert] 警告: %v", err)
	}

	// 初始更新
	updateSmartDNS()

	// 定时更新
	interval, _ := time.ParseDuration(cfg.Interval)
	if interval < time.Second {
		interval = 30 * time.Second
	}
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for range ticker.C {
			updateSmartDNS()
		}
	}()

	// 启动代理
	if cfg.SmartDNSIP != "" {
		log.Printf("[Proxy] HTTP 监听 :80")
		go http.ListenAndServe(":80", proxyHandler())
	}
	if cfg.Cert != "" && cfg.Key != "" {
		log.Printf("[Proxy] HTTPS 监听 :443")
		go http.ListenAndServeTLS(":443", cfg.Cert, cfg.Key, proxyHandler())
	}

	log.Printf("=== 全部服务已启动 ===")
	select {}
}

func proxyHandler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		host := strings.Split(r.Host, ":")[0]
		host = strings.TrimSpace(host)
		target := matchRule(host)

		if target == "" {
			// DIRECT 模式：查真实 IP
			log.Printf("[Proxy] DIRECT 模式: %s (TLS=%v)", host, r.TLS != nil)
			resolver := &net.Resolver{
				PreferGo: true,
				Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
					d := net.Dialer{Timeout: 5 * time.Second}
					return d.DialContext(ctx, "udp", "8.8.8.8:53")
				},
			}
			addrs, err := resolver.LookupHost(context.Background(), host)
			if err != nil || len(addrs) == 0 {
				http.Error(w, "503 DNS lookup failed", http.StatusServiceUnavailable)
				log.Printf("[Proxy] ❌ DNS 查询失败: %s → %v", host, err)
				return
			}
			realIP := addrs[0]
			port := "80"
			if r.TLS != nil {
				port = "443"
			}
			target = fmt.Sprintf("http://%s:%s", realIP, port)
			if r.TLS != nil {
				target = fmt.Sprintf("https://%s:%s", realIP, port)
			}
			log.Printf("[Proxy] ✅ DIRECT %s → %s", host, target)
		} else {
			// 规则匹配：根据客户端协议调整 target
			u, _ := url.Parse(target)
			if r.TLS == nil && u.Scheme == "https" {
				// 客户端 HTTP，但目标是 https → 改成 http
				u.Scheme = "http"
				h, p, _ := net.SplitHostPort(u.Host)
				if p == "443" || p == "" {
					u.Host = h + ":80"
				}
			}
			target = u.String()
			log.Printf("[Proxy] 规则匹配 %s → %s (客户端%s)", host, target, map[bool]string{true: "HTTPS", false: "HTTP"}[r.TLS != nil])
		}

		proxy := &httputil.ReverseProxy{
			Director: func(req *http.Request) {
				u, _ := url.Parse(target)
				req.URL.Scheme = u.Scheme
				req.URL.Host = u.Host
				req.Host = r.Host
				req.Header.Set("Host", r.Host)
				req.Header.Set("X-Forwarded-For", r.RemoteAddr)
			},
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			},
		}
		proxy.ServeHTTP(w, r)
	}
}

func matchRule(host string) string {
	host = strings.Split(host, ":")[0]
	host = strings.TrimSpace(host)

	// 1. 检查通配符规则（如 *.46.190564326.xyz）
	for _, rule := range cfg.Rules {
		if strings.HasPrefix(rule.Host, "*.") {
			suffix := rule.Host[1:] // 去掉 *，得到 .46.190564326.xyz
			if strings.HasSuffix(host, suffix) {
				if rule.Target != "" {
					return rule.Target
				}
				if cfg.LuckyAPITarget != "" {
					return cfg.LuckyAPITarget
				}
				return cfg.Proxy
			}
		}
	}

	// 2. 检查是否在 Lucky API 域名列表里
	if apiDomains[host] {
		if cfg.LuckyAPITarget != "" {
			return cfg.LuckyAPITarget
		}
		return cfg.Proxy
	}

	// 3. 检查是否在 config.json rules 里（完整域名）
	for _, rule := range cfg.Rules {
		if rule.Host == host && !strings.HasPrefix(rule.Host, "*.") {
			if rule.Target != "" && rule.Target != "DIRECT" {
				return rule.Target
			}
			break
		}
	}
	// 默认：返回空字符串（走 DIRECT 模式）
	return ""
}

func loadConfig(path string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, &cfg)
}

func updateSmartDNS() {
	changed := false

	// 1. 从 Lucky API 加载（过滤 46.）
	if loadFromAPI() {
		changed = true
	}

	// 2. 从 config.json rules 加载（不过滤）
	if loadRulesFromConf() {
		changed = true
	}

	if changed {
		mergeAndReload()
	}
}

func loadFromAPI() bool {
	domains, err := fetchDomains()
	if err != nil {
		log.Printf("[SmartDNS] 获取 Lucky 域名失败: %v", err)
		return false
	}
	if mapsEqual(domains, apiDomains) {
		return false
	}
	apiDomains = domains
	log.Printf("[SmartDNS] API 域名更新: %d 个", len(apiDomains))
	return true
}

func loadRulesFromConf() bool {
	// 快速检查：文件未改动就返回
	if cfg.ConfigFile != "" {
		info, err := os.Stat(cfg.ConfigFile)
		if err == nil && !info.ModTime().After(configModTime) {
			return false
		}
		configModTime = time.Now()
	}

	fileNew := make(map[string]bool)
	for _, rule := range cfg.Rules {
		if rule.Host != "" {
			fileNew[rule.Host] = true
		}
	}
	if mapsEqual(fileNew, fileDomains) {
		return false
	}
	fileDomains = fileNew
	log.Printf("[SmartDNS] Rules 域名更新: %d 个", len(fileDomains))
	return true
}

func mergeAndReload() {
	// 合并 apiDomains + fileDomains → allowedDomains
	allowedDomains = make(map[string]bool)
	for d := range apiDomains {
		allowedDomains[d] = true
	}
	for d := range fileDomains {
		allowedDomains[d] = true
	}

	// 生成 SmartDNS 配置
	lines := []string{
		fmt.Sprintf("# Auto-generated by lucky2smartdns @ %s", time.Now().Format("2006-01-02 15:04:05")),
		fmt.Sprintf("# Total: %d domains", len(allowedDomains)),
		fmt.Sprintf("# Target IP: %s", cfg.SmartDNSIP),
		"",
	}
	for d := range allowedDomains {
		lines = append(lines, fmt.Sprintf("address /%s/%s", d, cfg.SmartDNSIP))
	}
	content := strings.Join(lines, "\n")

	// 写到 /tmp/lucky-domains.conf
	if err := os.WriteFile(cfg.SmartDSConf, []byte(content), 0644); err != nil {
		log.Printf("[SmartDNS] 写文件失败: %v", err)
		return
	}

	// reload SmartDNS
	reloadSmartDNS()
}

func reloadSmartDNS() {
	cmd := exec.Command("sh", "-c", cfg.SmartDSReload)
	if err := cmd.Run(); err != nil {
		log.Printf("[SmartDNS] reload 失败: %v", err)
	} else {
		log.Printf("[SmartDNS] ✅ 配置已更新并 reload")
	}
}

func fetchDomains() (map[string]bool, error) {
	result := make(map[string]bool)
	// TODO: 实际调用 Lucky API
	return result, nil
}

func mapsEqual(a, b map[string]bool) bool {
	if len(a) != len(b) {
		return false
	}
	for k := range a {
		if !b[k] {
			return false
		}
	}
	return true
}

func loadCert() error {
	var err error
	cert, err = tls.LoadX509KeyPair(cfg.Cert, cfg.Key)
	return err
}