package main

import (
	"crypto/md5"
	"crypto/subtle"
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
)

const (
	downloadDir   = "./chrome_versions"
	checkInterval = 24 * time.Hour
	maxRetries    = 3
	updateURL     = "https://tools.google.com/service/update2"
)

var (
	authToken    string
	authEnabled  bool
	port         string
	keepVersions int
)

type Version struct {
	Filename string
	Size     int64
	Time     time.Time
}

type UpdateResponse struct {
	XMLName  xml.Name `xml:"response"`
	Protocol string   `xml:"protocol,attr"`
	App      struct {
		Status      string `xml:"status,attr"`
		UpdateCheck struct {
			Status string `xml:"status,attr"`
			URLs   struct {
				URL []struct {
					Codebase string `xml:"codebase,attr"`
				} `xml:"url"`
			} `xml:"urls"`
			Manifest struct {
				Version string `xml:"version,attr"`
				Packages struct {
					Package struct {
						Name       string `xml:"name,attr"`
						Required   bool   `xml:"required,attr"`
						Size       int64  `xml:"size,attr"`
						Hash       string `xml:"hash,attr"`
						HashSHA256 string `xml:"hash_sha256,attr"`
					} `xml:"package"`
				} `xml:"packages"`
			} `xml:"manifest"`
		} `xml:"updatecheck"`
	} `xml:"app"`
}

func init() {
	authToken = getEnv("AUTH_TOKEN", "")
	authEnabled = getEnv("AUTH_ENABLED", "false") == "true" || authToken != ""
	port = getEnv("PORT", "8080")
	keepVersions, _ = strconv.Atoi(getEnv("KEEP_VERSIONS", "3"))
}

func getEnv(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

func main() {
	os.MkdirAll(downloadDir, 0755)
	log.Printf("Auth enabled: %v, Port: %s, Keep versions: %d", authEnabled, port, keepVersions)

	log.Println("Performing initial check...")
	checkAndDownload()

	go monitor()

	http.HandleFunc("/", authMiddleware(serveIndex))
	http.Handle("/download/", authMiddleware(http.StripPrefix("/download/", http.FileServer(http.Dir(downloadDir))).ServeHTTP))

	log.Printf("Server started on :%s", port)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}

func monitor() {
	for {
		delay := time.Duration(rand.Intn(3600)) * time.Second
		time.Sleep(delay)

		checkAndDownload()
		time.Sleep(checkInterval)
	}
}

func getLatestVersionInfo() (version, downloadURL, sha256 string, err error) {
	reqBody := `<?xml version="1.0" encoding="UTF-8"?>
<request protocol="3.0" updater="Omaha" updaterversion="1.3.36.372" shell_version="1.3.36.352" ismachine="0" sessionid="{11111111-1111-1111-1111-111111111111}" installsource="taggedmi" requestid="{11111111-1111-1111-1111-111111111111}" dedup="cr" domainjoined="0">
<hw physmemory="16" sse="1" sse2="1" sse3="1" ssse3="1" sse41="1" sse42="1" avx="1"/>
<os platform="win" version="10.0.26100.1742" arch="x64"/>
<app version="" appid="{8A69D345-D564-463C-AFF1-A69D9E530F96}" ap="x64-stable">
<updatecheck/>
<data name="install" index="empty"/>
</app>
</request>`

	resp, err := http.Post(updateURL, "application/xml", strings.NewReader(reqBody))
	if err != nil {
		return "", "", "", err
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)

	var updateResp UpdateResponse
	if err := xml.Unmarshal(body, &updateResp); err != nil {
		return "", "", "", err
	}

	version = updateResp.App.UpdateCheck.Manifest.Version
	sha256 = updateResp.App.UpdateCheck.Manifest.Packages.Package.HashSHA256
	packageName := updateResp.App.UpdateCheck.Manifest.Packages.Package.Name

	if len(updateResp.App.UpdateCheck.URLs.URL) > 0 {
		for _, url := range updateResp.App.UpdateCheck.URLs.URL {
			if strings.HasPrefix(url.Codebase, "https://dl.google.com") {
				downloadURL = url.Codebase + packageName
				break
			}
		}
		if downloadURL == "" {
			downloadURL = updateResp.App.UpdateCheck.URLs.URL[0].Codebase + packageName
		}
	}

	return version, downloadURL, sha256, nil
}

func checkAndDownload() {
	version, downloadURL, sha256, err := getLatestVersionInfo()
	if err != nil {
		log.Printf("Failed to get version info: %v", err)
		return
	}

	log.Printf("Latest Chrome version: %s", version)
	log.Printf("Download URL: %s", downloadURL)
	log.Printf("SHA256: %s", sha256)

	if downloadURL == "" {
		log.Println("No download URL found")
		return
	}

	// 检查是否已存在该版本
	files, _ := os.ReadDir(downloadDir)
	for _, f := range files {
		if strings.HasPrefix(f.Name(), "chrome_"+version+"_") {
			log.Printf("Version %s already exists, skipping download", version)
			return
		}
	}

	var data []byte
	client := &http.Client{}

	for i := 0; i < maxRetries; i++ {
		resp, err := client.Get(downloadURL)
		if err != nil {
			log.Printf("Attempt %d failed: %v", i+1, err)
			time.Sleep(time.Duration(i+1) * 10 * time.Second)
			continue
		}

		data, err = io.ReadAll(resp.Body)
		resp.Body.Close()

		if err == nil && len(data) > 10000000 {
			break
		}

		log.Printf("Attempt %d: invalid response (size: %d bytes)", i+1, len(data))
		time.Sleep(time.Duration(i+1) * 10 * time.Second)
	}

	if len(data) < 10000000 {
		log.Println("All download attempts failed")
		return
	}

	hash := fmt.Sprintf("%x", md5.Sum(data))
	filename := fmt.Sprintf("chrome_%s_%s.exe", version, hash[:8])
	filepath := filepath.Join(downloadDir, filename)

	if err := os.WriteFile(filepath, data, 0644); err != nil {
		log.Printf("Error saving file: %v", err)
		return
	}

	log.Printf("Downloaded: %s (%.2f MB)", filename, float64(len(data))/1024/1024)
	cleanupOldVersions()
}

func cleanupOldVersions() {
	files, _ := os.ReadDir(downloadDir)
	if len(files) <= keepVersions {
		return
	}

	var versions []Version
	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".exe") {
			info, _ := f.Info()
			versions = append(versions, Version{
				Filename: f.Name(),
				Time:     info.ModTime(),
				Size:     info.Size(),
			})
		}
	}

	sort.Slice(versions, func(i, j int) bool {
		return versions[i].Time.After(versions[j].Time)
	})

	for i := keepVersions; i < len(versions); i++ {
		os.Remove(filepath.Join(downloadDir, versions[i].Filename))
		log.Printf("Removed old version: %s", versions[i].Filename)
	}
}

func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if authEnabled {
			token := r.URL.Query().Get("token")
			if subtle.ConstantTimeCompare([]byte(token), []byte(authToken)) != 1 {
				http.Error(w, "Unauthorized", http.StatusUnauthorized)
				return
			}
		}
		next(w, r)
	}
}

func serveIndex(w http.ResponseWriter, r *http.Request) {
	files, _ := os.ReadDir(downloadDir)
	var versions []Version

	for _, f := range files {
		if strings.HasSuffix(f.Name(), ".exe") {
			info, _ := f.Info()
			versions = append(versions, Version{
				Filename: f.Name(),
				Time:     info.ModTime(),
				Size:     info.Size(),
			})
		}
	}

	sort.Slice(versions, func(i, j int) bool {
		return versions[i].Time.After(versions[j].Time)
	})

	token := r.URL.Query().Get("token")
	tmpl := template.Must(template.New("index").Parse(`
<!DOCTYPE html>
<html>
<head><title>Chrome Offline Installer</title></head>
<body>
<h1>Chrome Offline Versions</h1>
<ul>
{{range .Versions}}
<li><a href="/download/{{.Filename}}{{if $.Token}}?token={{$.Token}}{{end}}">{{.Filename}}</a> ({{printf "%.2f" .SizeMB}} MB, {{.Time.Format "2006-01-02 15:04"}})</li>
{{end}}
</ul>
</body>
</html>
`))

	type VersionDisplay struct {
		Filename string
		SizeMB   float64
		Time     time.Time
	}

	var displayVersions []VersionDisplay
	for _, v := range versions {
		displayVersions = append(displayVersions, VersionDisplay{
			Filename: v.Filename,
			SizeMB:   float64(v.Size) / 1024 / 1024,
			Time:     v.Time,
		})
	}

	tmpl.Execute(w, struct {
		Versions []VersionDisplay
		Token    string
	}{displayVersions, token})
}