// main.go
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/sirupsen/logrus"
)

// Version 用于嵌入构建版本号
var Version = "dev"

// Config 定义配置结构体
type Config struct {
	ListenAddress string
	Port          int
	LogLevel      string
	DisguiseURL   string
}

var config Config

var client = &http.Client{
	CheckRedirect: func(req *http.Request, via []*http.Request) error {
		// 不跟随重定向，让 Docker 客户端自己处理
		return http.ErrUseLastResponse
	},
	Timeout: 30 * time.Second,
	Transport: &http.Transport{
		DisableKeepAlives:     false,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	},
}

type CustomFormatter struct {
	logrus.TextFormatter
}

func (f *CustomFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	timestamp := entry.Time.Format("2006-01-02 15:04:05.000")

	var levelColor string
	switch entry.Level {
	case logrus.DebugLevel:
		levelColor = "\033[36m"
	case logrus.InfoLevel:
		levelColor = "\033[32m"
	case logrus.WarnLevel:
		levelColor = "\033[33m"
	case logrus.ErrorLevel:
		levelColor = "\033[31m"
	case logrus.FatalLevel, logrus.PanicLevel:
		levelColor = "\033[35m"
	}

	resetColor := "\033[0m"
	logMessage := fmt.Sprintf("%s %s[%s]%s %s\n",
		timestamp,
		levelColor,
		strings.ToUpper(entry.Level.String()),
		resetColor,
		entry.Message)

	return []byte(logMessage), nil
}

func init() {
	logrus.SetFormatter(&CustomFormatter{
		TextFormatter: logrus.TextFormatter{
			DisableColors:   false,
			FullTimestamp:   true,
			TimestampFormat: "2006-01-02 15:04:05.000",
		},
	})
}

func preprocessArgs() {
	alias := map[string]string{
		"--listen":    "-l",
		"--port":      "-p",
		"--log-level": "-ll",
		"--disguise":  "-w",
	}

	newArgs := make([]string, 0, len(os.Args))
	newArgs = append(newArgs, os.Args[0])

	for _, arg := range os.Args[1:] {
		if strings.HasPrefix(arg, "--") && strings.Contains(arg, "=") {
			parts := strings.SplitN(arg, "=", 2)
			if short, ok := alias[parts[0]]; ok {
				arg = short + "=" + parts[1]
			}
		} else if short, ok := alias[arg]; ok {
			arg = short
		}
		newArgs = append(newArgs, arg)
	}

	if len(newArgs) > 1 {
		os.Args = newArgs
	}
}

func usage() {
	const helpText = `HubP - Docker Hub 代理服务器

参数说明:
    -l, --listen       监听地址 (默认: 0.0.0.0)
    -p, --port         监听端口 (默认: 18184)
    -ll, --log-level   日志级别: debug/info/warn/error (默认: info)
    -w, --disguise     伪装网站 URL (默认: onlinealarmkur.com)

示例:
    ./HubP -l 0.0.0.0 -p 18184 -ll debug -w www.bing.com
    ./HubP --listen=0.0.0.0 --port=18184 --log-level=debug --disguise=www.bing.com`

	fmt.Fprintf(os.Stderr, "%s\n", helpText)
}

func validateConfig() error {
	if config.Port < 1 || config.Port > 65535 {
		return fmt.Errorf("无效的端口号: %d", config.Port)
	}
	if config.DisguiseURL == "" {
		return fmt.Errorf("伪装网站 URL 不能为空")
	}
	return nil
}

func main() {
	preprocessArgs()
	flag.Usage = usage

	defaultListenAddress := getEnv("HUBP_LISTEN", "0.0.0.0")
	defaultPort := getEnvAsInt("HUBP_PORT", 18184)
	defaultLogLevel := getEnv("HUBP_LOG_LEVEL", "debug")
	defaultDisguiseURL := getEnv("HUBP_DISGUISE", "onlinealarmkur.com")

	flag.StringVar(&config.ListenAddress, "l", defaultListenAddress, "监听地址")
	flag.IntVar(&config.Port, "p", defaultPort, "监听端口")
	flag.StringVar(&config.LogLevel, "ll", defaultLogLevel, "日志级别")
	flag.StringVar(&config.DisguiseURL, "w", defaultDisguiseURL, "伪装网站 URL")

	if err := flag.CommandLine.Parse(os.Args[1:]); err != nil {
		logrus.Fatal("解析命令行参数失败：", err)
	}

	level, err := logrus.ParseLevel(config.LogLevel)
	if err != nil {
		logrus.Warnf("无效的日志级别 '%s'，使用默认级别 'info'", config.LogLevel)
		level = logrus.InfoLevel
	}
	logrus.SetLevel(level)

	if err := validateConfig(); err != nil {
		logrus.Fatal("配置验证失败: ", err)
	}

	printStartupInfo()

	addr := fmt.Sprintf("%s:%d", config.ListenAddress, config.Port)
	http.HandleFunc("/", handleRequest)

	server := &http.Server{
		Addr:    addr,
		Handler: http.DefaultServeMux,
	}

	go func() {
		logrus.Info("服务启动成功")
		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logrus.Fatal("服务启动失败: ", err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	logrus.Info("正在关闭服务器...")
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := server.Shutdown(ctx); err != nil {
		logrus.Error("服务器强制关闭: ", err)
	}
	logrus.Info("服务器已关闭")
}

func printStartupInfo() {
	const blue = "\033[34m"
	const green = "\033[32m"
	const reset = "\033[0m"

	fmt.Println(blue + "\n╔════════════════════════════════════════════════════════════╗" + reset)
	fmt.Println(blue + "║" + green + "               HubP Docker Hub 代理服务器               " + blue + "║" + reset)
	fmt.Printf(blue+"║"+green+"               版本: %-33s"+blue+"║\n"+reset, Version)
	fmt.Println(blue + "╠════════════════════════════════════════════════════════════╣" + reset)
	fmt.Printf(blue+"║"+reset+" 监听地址: %-43s"+blue+"║\n"+reset, config.ListenAddress)
	fmt.Printf(blue+"║"+reset+" 监听端口: %-43d"+blue+"║\n"+reset, config.Port)
	fmt.Printf(blue+"║"+reset+" 日志级别: %-43s"+blue+"║\n"+reset, config.LogLevel)
	fmt.Printf(blue+"║"+reset+" 伪装网站: %-43s"+blue+"║\n"+reset, config.DisguiseURL)
	fmt.Println(blue + "╚════════════════════════════════════════════════════════════╝" + reset)
	fmt.Println()
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path

	// 健康检查
	if path == "/health" || path == "/healthz" {
		w.WriteHeader(http.StatusOK)
		if _, err := w.Write([]byte("OK")); err != nil {
			logrus.Errorf("健康检查响应失败: %v", err)
		}
		return
	}

	if logrus.IsLevelEnabled(logrus.DebugLevel) {
		var routeTag string
		if strings.HasPrefix(path, "/v2/") {
			routeTag = "[Docker]"
		} else if strings.HasPrefix(path, "/auth/") {
			routeTag = "[认证]"
		} else if strings.HasPrefix(path, "/production-cloudflare/") {
			routeTag = "[CF]"
		} else {
			routeTag = "[伪装]"
		}

		logrus.Debugf("%s 请求: [%s %s] 来自 %s",
			routeTag, r.Method, r.URL.String(), r.RemoteAddr)
	}

	if strings.HasPrefix(path, "/v2/") {
		handleRegistryRequest(w, r)
	} else if strings.HasPrefix(path, "/auth/") {
		handleAuthRequest(w, r)
	} else if strings.HasPrefix(path, "/production-cloudflare/") {
		handleCloudflareRequest(w, r)
	} else {
		handleDisguise(w, r)
	}
}

func getCleanHost(r *http.Request) string {
	host := r.Host
	if idx := strings.Index(host, ":"); idx != -1 {
		return host[:idx]
	}
	return host
}

func handleRegistryRequest(w http.ResponseWriter, r *http.Request) {
	const targetHost = "registry-1.docker.io"

	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()

	path := strings.TrimPrefix(r.URL.Path, "/v2/")

	url := &url.URL{
		Scheme:   "https",
		Host:     targetHost,
		Path:     "/v2/" + path,
		RawQuery: r.URL.RawQuery,
	}

	headers := copyHeaders(r.Header)
	headers.Set("Host", targetHost)

	logrus.Debugf("Docker镜像: 转发请求至 %s", url.String())

	resp, err := sendRequestWithContext(ctx, r.Method, url.String(), headers, r.Body)
	if err != nil {
		logrus.Errorf("Docker镜像: 请求失败 - %v", err)
		if logrus.IsLevelEnabled(logrus.DebugLevel) {
			http.Error(w, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway)
		} else {
			http.Error(w, "服务暂时不可用", http.StatusBadGateway)
		}
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusUnauthorized {
		handleAuthChallenge(w, r, resp)
		return
	}

	respHeaders := copyHeaders(resp.Header)

	if respHeaders.Get("WWW-Authenticate") != "" {
		currentDomain := getCleanHost(r)
		respHeaders.Set("WWW-Authenticate",
			fmt.Sprintf(`Bearer realm="https://%s/auth/token", service="registry.docker.io"`, currentDomain))
	}

	for k, v := range respHeaders {
		for _, val := range v {
			w.Header().Add(k, val)
		}
	}
	w.WriteHeader(resp.StatusCode)

	written, err := io.Copy(w, resp.Body)
	if err != nil {
		logrus.Errorf("Docker镜像: 传输响应失败 - %v", err)
		return
	}

	if logrus.IsLevelEnabled(logrus.DebugLevel) {
		logrus.Debugf("Docker镜像: 响应完成 [状态: %d] [大小: %.2f KB]",
			resp.StatusCode, float64(written)/1024)
	}
}

func handleAuthRequest(w http.ResponseWriter, r *http.Request) {
	const targetHost = "auth.docker.io"

	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()

	path := strings.TrimPrefix(r.URL.Path, "/auth/")

	url := &url.URL{
		Scheme:   "https",
		Host:     targetHost,
		Path:     "/" + path,
		RawQuery: r.URL.RawQuery,
	}

	headers := copyHeaders(r.Header)
	headers.Set("Host", targetHost)

	logrus.Debugf("认证服务: 转发请求至 %s", url.String())

	resp, err := sendRequestWithContext(ctx, r.Method, url.String(), headers, r.Body)
	if err != nil {
		logrus.Errorf("认证服务: 请求失败 - %v", err)
		if logrus.IsLevelEnabled(logrus.DebugLevel) {
			http.Error(w, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway)
		} else {
			http.Error(w, "服务暂时不可用", http.StatusBadGateway)
		}
		return
	}
	defer resp.Body.Close()

	for k, v := range resp.Header {
		for _, val := range v {
			w.Header().Add(k, val)
		}
	}
	w.WriteHeader(resp.StatusCode)

	written, err := io.Copy(w, resp.Body)
	if err != nil {
		logrus.Errorf("认证服务: 传输响应失败 - %v", err)
		return
	}

	if logrus.IsLevelEnabled(logrus.DebugLevel) {
		logrus.Debugf("认证服务: 响应完成 [状态: %d] [大小: %.2f KB]",
			resp.StatusCode, float64(written)/1024)
	}
}

func handleCloudflareRequest(w http.ResponseWriter, r *http.Request) {
	const targetHost = "production.cloudflare.docker.com"

	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()

	path := strings.TrimPrefix(r.URL.Path, "/production-cloudflare/")

	url := &url.URL{
		Scheme:   "https",
		Host:     targetHost,
		Path:     "/" + path,
		RawQuery: r.URL.RawQuery,
	}

	headers := copyHeaders(r.Header)
	headers.Set("Host", targetHost)

	logrus.Debugf("Cloudflare: 转发请求至 %s", url.String())

	resp, err := sendRequestWithContext(ctx, r.Method, url.String(), headers, r.Body)
	if err != nil {
		logrus.Errorf("Cloudflare: 请求失败 - %v", err)
		if logrus.IsLevelEnabled(logrus.DebugLevel) {
			http.Error(w, fmt.Sprintf("代理错误: %v", err), http.StatusBadGateway)
		} else {
			http.Error(w, "服务暂时不可用", http.StatusBadGateway)
		}
		return
	}
	defer resp.Body.Close()

	for k, v := range resp.Header {
		for _, val := range v {
			w.Header().Add(k, val)
		}
	}
	w.WriteHeader(resp.StatusCode)

	written, err := io.Copy(w, resp.Body)
	if err != nil {
		logrus.Errorf("Cloudflare: 传输响应失败 - %v", err)
		return
	}

	if logrus.IsLevelEnabled(logrus.DebugLevel) {
		logrus.Debugf("Cloudflare: 响应完成 [状态: %d] [大小: %.2f KB]",
			resp.StatusCode, float64(written)/1024)
	}
}

func handleAuthChallenge(w http.ResponseWriter, r *http.Request, resp *http.Response) {
	for k, v := range resp.Header {
		for _, val := range v {
			w.Header().Add(k, val)
		}
	}

	if authHeader := w.Header().Get("WWW-Authenticate"); authHeader != "" {
		currentDomain := getCleanHost(r)
		w.Header().Set("WWW-Authenticate",
			fmt.Sprintf(`Bearer realm="https://%s/auth/token", service="registry.docker.io"`, currentDomain))
	}

	w.WriteHeader(resp.StatusCode)

	_, err := io.Copy(w, resp.Body)
	if err != nil {
		logrus.Errorf("认证响应传输失败: %v", err)
	}
}

func handleDisguise(w http.ResponseWriter, r *http.Request) {
	targetURL := &url.URL{
		Scheme:   "https",
		Host:     config.DisguiseURL,
		Path:     r.URL.Path,
		RawQuery: r.URL.RawQuery,
	}

	if logrus.IsLevelEnabled(logrus.DebugLevel) {
		logrus.Debugf("伪装页面: 转发请求至 %s", targetURL.String())
	}

	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()

	headers := copyHeaders(r.Header)
	headers.Del("Accept-Encoding")

	resp, err := sendRequestWithContext(ctx, r.Method, targetURL.String(), headers, r.Body)
	if err != nil {
		logrus.Errorf("伪装页面: 请求失败 - %v", err)
		http.Error(w, "服务器错误", http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	for k, v := range resp.Header {
		for _, val := range v {
			w.Header().Add(k, val)
		}
	}
	w.WriteHeader(resp.StatusCode)

	written, err := io.Copy(w, resp.Body)
	if err != nil {
		logrus.Errorf("伪装页面: 传输响应失败 - %v", err)
		return
	}

	if logrus.IsLevelEnabled(logrus.DebugLevel) {
		logrus.Debugf("伪装页面: 响应完成 [状态: %d] [大小: %.2f KB]",
			resp.StatusCode, float64(written)/1024)
	}
}

func sendRequestWithContext(ctx context.Context, method, url string, headers http.Header, body io.ReadCloser) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, method, url, body)
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %v", err)
	}

	req.Header = headers
	
	startTime := time.Now()
	resp, err := client.Do(req)

	if err == nil && logrus.IsLevelEnabled(logrus.DebugLevel) {
		duration := time.Since(startTime)
		logrus.Debugf("请求耗时: %.2f 秒 (%s)", duration.Seconds(), url)
	}

	return resp, err
}

func copyHeaders(src http.Header) http.Header {
	dst := make(http.Header)
	for key, values := range src {
		dst[key] = append([]string(nil), values...)
	}
	return dst
}

func getEnv(key, defaultValue string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return defaultValue
}

func getEnvAsInt(key string, defaultValue int) int {
	if valueStr, exists := os.LookupEnv(key); exists {
		if value, err := strconv.Atoi(valueStr); err == nil {
			return value
		}
	}
	return defaultValue
}