// main.go
package main

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

	"github.com/sirupsen/logrus"
)

var Version = "dev"

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 {
		return http.ErrUseLastResponse
	},
	Timeout: 30 * time.Second,
	Transport: &http.Transport{
		DisableKeepAlives:     false,
		MaxIdleConns:          100,
		MaxIdleConnsPerHost:   10,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	},
}

// Token 缓存
type TokenCache struct {
	mu    sync.RWMutex
	cache map[string]*CachedToken
}

type CachedToken struct {
	Token     string
	ExpiresAt time.Time
}

var tokenCache = &TokenCache{
	cache: make(map[string]*CachedToken),
}

func (tc *TokenCache) Get(key string) (string, bool) {
	tc.mu.RLock()
	defer tc.mu.RUnlock()
	
	if cached, ok := tc.cache[key]; ok {
		if time.Now().Before(cached.ExpiresAt) {
			return cached.Token, true
		}
		delete(tc.cache, key)
	}
	return "", false
}

func (tc *TokenCache) Set(key, token string, ttl time.Duration) {
	tc.mu.Lock()
	defer tc.mu.Unlock()
	
	tc.cache[key] = &CachedToken{
		Token:     token,
		ExpiresAt: time.Now().Add(ttl),
	}
}

// Manifest 缓存
type ManifestCache struct {
	mu    sync.RWMutex
	cache map[string]*CachedManifest
}

type CachedManifest struct {
	Data      []byte
	Headers   http.Header
	ExpiresAt time.Time
}

var manifestCache = &ManifestCache{
	cache: make(map[string]*CachedManifest),
}

func (mc *ManifestCache) Get(key string) (*CachedManifest, bool) {
	mc.mu.RLock()
	defer mc.mu.RUnlock()
	
	if cached, ok := mc.cache[key]; ok {
		if time.Now().Before(cached.ExpiresAt) {
			return cached, true
		}
		delete(mc.cache, key)
	}
	return nil, false
}

func (mc *ManifestCache) Set(key string, data []byte, headers http.Header, ttl time.Duration) {
	mc.mu.Lock()
	defer mc.mu.Unlock()
	
	mc.cache[key] = &CachedManifest{
		Data:      data,
		Headers:   headers,
		ExpiresAt: time.Now().Add(ttl),
	}
}

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"
	return []byte(fmt.Sprintf("%s %s[%s]%s %s\n",
		timestamp, levelColor, strings.ToUpper(entry.Level.String()), resetColor, entry.Message)), 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() {
	fmt.Fprintf(os.Stderr, `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
`)
}

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

	flag.StringVar(&config.ListenAddress, "l", getEnv("HUBP_LISTEN", "0.0.0.0"), "监听地址")
	flag.IntVar(&config.Port, "p", getEnvAsInt("HUBP_PORT", 18184), "监听端口")
	flag.StringVar(&config.LogLevel, "ll", getEnv("HUBP_LOG_LEVEL", "debug"), "日志级别")
	flag.StringVar(&config.DisguiseURL, "w", getEnv("HUBP_DISGUISE", "onlinealarmkur.com"), "伪装网站 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, green, reset = "\033[34m", "\033[32m", "\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)
		w.Write([]byte("OK"))
		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"
	path := strings.TrimPrefix(r.URL.Path, "/v2/")

	// Manifest 缓存检查
	if r.Method == http.MethodGet && strings.Contains(path, "/manifests/") {
		cacheKey := fmt.Sprintf("manifest:%s", path)
		if cached, ok := manifestCache.Get(cacheKey); ok {
			logrus.Debugf("Docker镜像: 使用缓存的 manifest")
			for k, v := range cached.Headers {
				for _, val := range v {
					w.Header().Add(k, val)
				}
			}
			w.WriteHeader(http.StatusOK)
			w.Write(cached.Data)
			return
		}
	}

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

	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)
		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)

	// 缓存 manifest
	if resp.StatusCode == http.StatusOK && r.Method == http.MethodGet && strings.Contains(path, "/manifests/") {
		data, err := io.ReadAll(resp.Body)
		if err == nil {
			cacheKey := fmt.Sprintf("manifest:%s", path)
			ttl := 10 * time.Minute
			if strings.Contains(path, "sha256:") {
				ttl = 1 * time.Hour
			}
			manifestCache.Set(cacheKey, data, respHeaders, ttl)
			w.Write(data)
			logrus.Debugf("Docker镜像: manifest 已缓存 [大小: %.2f KB]", float64(len(data))/1024)
			return
		}
	}

	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"

	// Token 缓存检查
	cacheKey := r.URL.RawQuery
	if token, ok := tokenCache.Get(cacheKey); ok {
		logrus.Debugf("认证服务: 使用缓存的 token")
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(token))
		return
	}

	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)
		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)

	// 缓存 token
	if resp.StatusCode == http.StatusOK {
		data, err := io.ReadAll(resp.Body)
		if err == nil {
			var tokenResp map[string]interface{}
			if json.Unmarshal(data, &tokenResp) == nil {
				ttl := 5 * time.Minute
				if expiresIn, ok := tokenResp["expires_in"].(float64); ok {
					ttl = time.Duration(expiresIn*0.9) * time.Second
				}
				tokenCache.Set(cacheKey, string(data), ttl)
				logrus.Debugf("认证服务: token 已缓存 [TTL: %v]", ttl)
			}
			w.Write(data)
			return
		}
	}

	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)
		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)
	io.Copy(w, resp.Body)
}

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
}