// 源站压力探针：连续快速请求源站，看它是否也限流。
// 同时校验源站图和 Photon 图的字节数是否一致（确认画质没变）。
//
//	go run ./tools/probeorigin -load <Photon图片URL> [次数]
package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

func main() {
	load := false
	ref := false
	var url string
	n := 20
	for i := 1; i < len(os.Args); i++ {
		if os.Args[i] == "-load" {
			load = true
			continue
		}
		if os.Args[i] == "-ref" {
			ref = true
			continue
		}
		if url == "" {
			url = os.Args[i]
		} else if v, err := strconv.Atoi(os.Args[i]); err == nil {
			n = v
		}
	}
	if url == "" {
		fmt.Fprintln(os.Stderr, "用法: go run ./tools/probeorigin [-load|-ref] <Photon或源站URL> [次数]")
		os.Exit(1)
	}
	origin, ok := stripPhoton(url)
	if !ok {
		origin = url // 已经是源站直链
	}

	// 防盗链探测：源站到底认不认 Referer？
	//
	// 风险点：浏览器走 Photon 而非源站，说明源站可能有防盗链、只是被
	// Photon 绕开了。若是这样，直连就必须带对 Referer。
	if ref {
		fmt.Printf("源站: %s\n\n", origin)
		for _, c := range []struct{ name, referer string }{
			{"裸请求(不带头)", ""},
			{"本域 Referer", "https://migmg.ssuu.uk/content/09/x.html"},
			{"跨站 Referer", "https://example.com/some-page"},
			{"Photon 自己", "https://i0.wp.com/"},
		} {
			status, b, ct, err := fetchRef(origin, c.referer, c.referer == "")
			if err != nil {
				fmt.Printf("%-16s → 错误: %v\n", c.name, err)
				continue
			}
			mark := "✓"
			if status != 200 || b < 1000 {
				mark = "✗"
			}
			fmt.Printf("%-16s → HTTP %d  %d 字节  type=%s  %s\n", c.name, status, b, ct, mark)
			time.Sleep(700 * time.Millisecond)
		}
		return
	}

	if !load {
		fmt.Printf("源站: %s\n", origin)
		status, b, ct, err := fetch(origin)
		fmt.Printf("HTTP %d  %d 字节  type=%s  err=%v\n", status, b, ct, err)
		return
	}

	// 并发压测源站
	fmt.Printf("压测源站 ×%d 并发...\n", n)
	var ok200, fail int64
	var mu sync.Mutex
	lat := []time.Duration{}

	start := time.Now()
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			t := time.Now()
			status, b, _, err := fetch(origin)
			d := time.Since(t)
			mu.Lock()
			lat = append(lat, d)
			mu.Unlock()
			if err == nil && status == 200 && b > 1000 {
				atomic.AddInt64(&ok200, 1)
			} else {
				atomic.AddInt64(&fail, 1)
			}
		}()
	}
	wg.Wait()
	total := time.Since(start)

	fmt.Printf("成功 %d / 失败 %d   总耗时 %v\n", ok200, fail, total.Round(time.Millisecond))
	if len(lat) > 0 {
		var sum time.Duration
		for _, d := range lat {
			sum += d
		}
		fmt.Printf("平均延迟 %v  最快 %v  最慢 %v\n",
			(sum / time.Duration(len(lat))).Round(time.Millisecond),
			minDur(lat).Round(time.Millisecond),
			maxDur(lat).Round(time.Millisecond))
	}
}

func minDur(d []time.Duration) time.Duration {
	m := d[0]
	for _, x := range d[1:] {
		if x < m {
			m = x
		}
	}
	return m
}

func maxDur(d []time.Duration) time.Duration {
	m := d[0]
	for _, x := range d[1:] {
		if x > m {
			m = x
		}
	}
	return m
}

func stripPhoton(raw string) (string, bool) {
	low := strings.ToLower(raw)
	idx := strings.Index(low, "://i")
	if idx < 0 {
		return "", false
	}
	rest := low[idx+3:]
	dot := strings.Index(rest, ".wp.com/")
	if dot < 0 {
		return "", false
	}
	scheme := raw[:idx]
	origin := scheme + "://" + raw[idx+3+dot+len(".wp.com/"):]
	if i := strings.IndexAny(origin, "?#"); i >= 0 {
		origin = origin[:i]
	}
	return origin, true
}

func fetch(url string) (int, int64, string, error) {
	return fetchRef(url, "https://migmg.ssuu.uk/", false)
}

// fetchRef 带可控 Referer 的请求。bare=true 时连 UA 都不发，
// 用来验证「源站是否真的无防盗链」。
func fetchRef(url, referer string, bare bool) (int, int64, string, error) {
	tr := &http.Transport{ForceAttemptHTTP2: true}
	defer tr.CloseIdleConnections()
	client := &http.Client{Transport: tr, Timeout: 25 * time.Second}

	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return 0, 0, "", err
	}
	if !bare {
		req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36")
		req.Header.Set("Accept", "image/webp,image/*,*/*;q=0.8")
	}
	if referer != "" {
		req.Header.Set("Referer", referer)
	}

	resp, err := client.Do(req)
	if err != nil {
		return 0, 0, "", err
	}
	defer resp.Body.Close()
	b, _ := io.Copy(io.Discard, resp.Body)
	return resp.StatusCode, b, resp.Header.Get("Content-Type"), nil
}