// 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/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"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{
	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,
	},
}

var remoteOptions = []remote.Option{
	remote.WithAuth(authn.Anonymous),
	remote.WithTransport(client.Transport),
}

// 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
	MediaType string
	Digest    string
	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, mediaType, digest string, ttl time.Duration) {
	mc.mu.Lock()
	defer mc.mu.Unlock()
	mc.cache[key] = &CachedManifest{
		Data:      data,
		MediaType: mediaType,
		Digest:    digest,
		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) {
	path := strings.TrimPrefix(r.URL.Path, "/v2/")

	// /v2/ 端点
	if path == "" {
		w.Header().Set("Docker-Distribution-API-Version", "registry/2.0")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("{}"))
		return
	}

	imageName, apiType, reference := parseRegistryPath(path)
	if imageName == "" || apiType == "" {
		http.Error(w, "Invalid path format", http.StatusBadRequest)
		return
	}

	if !strings.Contains(imageName, "/") {
		imageName = "library/" + imageName
	}

	imageRef := fmt.Sprintf("registry-1.docker.io/%s", imageName)

	switch apiType {
	case "manifests":
		handleManifestRequest(w, r, imageRef, reference)
	case "blobs":
		handleBlobRequest(w, r, imageRef, reference)
	default:
		http.Error(w, "API endpoint not found", http.StatusNotFound)
	}
}

func parseRegistryPath(path string) (imageName, apiType, reference string) {
	if idx := strings.Index(path, "/manifests/"); idx != -1 {
		imageName = path[:idx]
		apiType = "manifests"
		reference = path[idx+len("/manifests/"):]
		return
	}
	if idx := strings.Index(path, "/blobs/"); idx != -1 {
		imageName = path[:idx]
		apiType = "blobs"
		reference = path[idx+len("/blobs/"):]
		return
	}
	return "", "", ""
}

func handleManifestRequest(w http.ResponseWriter, r *http.Request, imageRef, reference string) {
	cacheKey := fmt.Sprintf("%s@%s", imageRef, reference)

	// 检查缓存
	if r.Method == http.MethodGet {
		if cached, ok := manifestCache.Get(cacheKey); ok {
			logrus.Debugf("Docker镜像: 使用缓存的 manifest")
			w.Header().Set("Content-Type", cached.MediaType)
			w.Header().Set("Docker-Content-Digest", cached.Digest)
			w.Header().Set("Content-Length", fmt.Sprintf("%d", len(cached.Data)))
			w.WriteHeader(http.StatusOK)
			w.Write(cached.Data)
			return
		}
	}

	var ref name.Reference
	var err error

	if strings.HasPrefix(reference, "sha256:") {
		ref, err = name.NewDigest(fmt.Sprintf("%s@%s", imageRef, reference))
	} else {
		ref, err = name.NewTag(fmt.Sprintf("%s:%s", imageRef, reference))
	}

	if err != nil {
		logrus.Errorf("解析镜像引用失败: %v", err)
		http.Error(w, "Invalid reference", http.StatusBadRequest)
		return
	}

	if r.Method == http.MethodHead {
		desc, err := remote.Head(ref, remoteOptions...)
		if err != nil {
			logrus.Errorf("HEAD请求失败: %v", err)
			http.Error(w, "Manifest not found", http.StatusNotFound)
			return
		}

		w.Header().Set("Content-Type", string(desc.MediaType))
		w.Header().Set("Docker-Content-Digest", desc.Digest.String())
		w.Header().Set("Content-Length", fmt.Sprintf("%d", desc.Size))
		w.WriteHeader(http.StatusOK)
		logrus.Debugf("Docker镜像: HEAD 响应完成")
	} else {
		desc, err := remote.Get(ref, remoteOptions...)
		if err != nil {
			logrus.Errorf("GET请求失败: %v", err)
			http.Error(w, "Manifest not found", http.StatusNotFound)
			return
		}

		w.Header().Set("Content-Type", string(desc.MediaType))
		w.Header().Set("Docker-Content-Digest", desc.Digest.String())
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(desc.Manifest)))
		w.WriteHeader(http.StatusOK)
		w.Write(desc.Manifest)

		// 缓存 manifest
		ttl := 10 * time.Minute
		if strings.HasPrefix(reference, "sha256:") {
			ttl = 1 * time.Hour
		}
		manifestCache.Set(cacheKey, desc.Manifest, string(desc.MediaType), desc.Digest.String(), ttl)

		logrus.Debugf("Docker镜像: manifest 响应完成 [大小: %.2f KB]", float64(len(desc.Manifest))/1024)
	}
}

func handleBlobRequest(w http.ResponseWriter, r *http.Request, imageRef, digest string) {
	digestRef, err := name.NewDigest(fmt.Sprintf("%s@%s", imageRef, digest))
	if err != nil {
		logrus.Errorf("解析digest引用失败: %v", err)
		http.Error(w, "Invalid digest reference", http.StatusBadRequest)
		return
	}

	layer, err := remote.Layer(digestRef, remoteOptions...)
	if err != nil {
		logrus.Errorf("获取layer失败: %v", err)
		http.Error(w, "Layer not found", http.StatusNotFound)
		return
	}

	size, err := layer.Size()
	if err != nil {
		logrus.Errorf("获取layer大小失败: %v", err)
		http.Error(w, "Failed to get layer size", http.StatusInternalServerError)
		return
	}

	reader, err := layer.Compressed()
	if err != nil {
		logrus.Errorf("获取layer内容失败: %v", err)
		http.Error(w, "Failed to get layer content", http.StatusInternalServerError)
		return
	}
	defer reader.Close()

	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Length", fmt.Sprintf("%d", size))
	w.Header().Set("Docker-Content-Digest", digest)
	w.WriteHeader(http.StatusOK)

	written, err := io.Copy(w, reader)
	if err != nil {
		logrus.Errorf("传输layer失败: %v", err)
		return
	}

	logrus.Debugf("Docker镜像: blob 传输完成 [大小: %.2f MB]", float64(written)/(1024*1024))
}

func handleAuthRequest(w http.ResponseWriter, r *http.Request) {
	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
	}

	const targetHost = "auth.docker.io"
	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()

	path := strings.TrimPrefix(r.URL.Path, "/auth/")
	targetURL := &url.URL{
		Scheme:   "https",
		Host:     targetHost,
		Path:     "/" + path,
		RawQuery: r.URL.RawQuery,
	}

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

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

	resp, err := sendRequestWithContext(ctx, r.Method, targetURL.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)

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

	io.Copy(w, resp.Body)
}

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/")
	targetURL := &url.URL{
		Scheme:   "https",
		Host:     targetHost,
		Path:     "/" + path,
		RawQuery: r.URL.RawQuery,
	}

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

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

	resp, err := sendRequestWithContext(ctx, r.Method, targetURL.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)
	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,
	}

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

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
	return client.Do(req)
}

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
}