eh-downloader/internal/core/engine.go

package core

// 下载引擎：任务编排 / 并发控制 / 指数退避重试 / 断点续传
//
// 断点续传策略：以「文件系统为准」——
//
//	目标文件已存在且 size > 0 → 视为已完成，直接跳过。
//
// 因此即使任务状态 JSON 丢失或损坏，重跑也能正确续传。

import (
	"context"
	"errors"
	"fmt"
	"math"
	"math/rand"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"eh-downloader-wails/internal/eh"
)

// ErrCanceled 取消标记（用户主动停止不算失败）
var ErrCanceled = errors.New("已取消")

// retryableErrs 可重试的错误特征（对应 JS 版 RETRYABLE 集合）
var retryablePatterns = []string{
	"connection reset", "connection refused", "timeout", "i/o timeout",
	"broken pipe", "eof", "no such host", "tls handshake",
	"unexpected status", "server closed", "too many redirects",
	"429", "502", "503", "504", "socket hang up",
}

// EventSink 是引擎向 UI 发事件的方式（由 app 层注入）
type EventSink interface {
	TaskUpdate(t Task)
	Log(level, msg string)
	Done()
}

// Engine 下载引擎
type Engine struct {
	store  *Store
	client *eh.Client
	sink   EventSink

	mu              sync.Mutex
	galleryWorkers  int
	pageConcurrency int
	running         bool
	cancel          context.CancelFunc
	queue           []Task
}

// NewEngine 建引擎
func NewEngine(store *Store, sink EventSink, settings Settings) *Engine {
	e := &Engine{
		store:           store,
		sink:            sink,
		client:          eh.NewClient(time.Duration(settings.Timeout)*time.Millisecond, settings.Proxy, settings.ProxyMode, settings.Xres),
		galleryWorkers:  clampInt(settings.GalleryConcurrency, 1, 20, 5),
		pageConcurrency: clampInt(settings.PageConcurrency, 1, 20, 8),
	}
	e.client.SetRequestInterval(time.Duration(settings.RequestIntervalMs) * time.Millisecond)
	return e
}

// clampInt 把值夹到 [lo, hi]。0 或负数视为「没设置」，用默认值。
func clampInt(v, lo, hi, def int) int {
	if v <= 0 {
		return def
	}
	if v < lo {
		return lo
	}
	if v > hi {
		return hi
	}
	return v
}

// Client 暴露给 app 层改配置用
func (e *Engine) Client() *eh.Client { return e.client }

// SetConcurrency 运行时改并发
func (e *Engine) SetConcurrency(n int) {
	e.mu.Lock()
	e.galleryWorkers = clampInt(n, 1, 20, 5)
	e.mu.Unlock()
}

// SetGalleryConcurrency 改画廊并发
func (e *Engine) SetGalleryConcurrency(n int) { e.SetConcurrency(n) }

// SetPageConcurrency 改单画廊页并发
func (e *Engine) SetPageConcurrency(n int) {
	e.mu.Lock()
	e.pageConcurrency = clampInt(n, 1, 20, 8)
	e.mu.Unlock()
}

// SetRequestInterval 改全局请求间隔（毫秒）
func (e *Engine) SetRequestInterval(ms int) {
	if ms < 0 {
		ms = 0
	}
	e.client.SetRequestInterval(time.Duration(ms) * time.Millisecond)
}

// GalleryConcurrency 当前画廊并发
func (e *Engine) GalleryConcurrency() int {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.galleryWorkers
}

// PageConcurrency 当前页并发
func (e *Engine) PageConcurrency() int {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.pageConcurrency
}

// emitTask 把任务状态推给 UI
func (e *Engine) emitTask(t Task) {
	if e.sink != nil {
		e.sink.TaskUpdate(t)
	}
}

func (e *Engine) IsRunning() bool {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.running
}

func (e *Engine) QueueLength() int {
	e.mu.Lock()
	defer e.mu.Unlock()
	return len(e.queue)
}

// Stop 停止：取消 context，立刻打断在途请求
func (e *Engine) Stop() {
	e.mu.Lock()
	if e.cancel != nil {
		e.cancel()
		e.cancel = nil
	}
	e.mu.Unlock()
}

// Enqueue 把任务加入队列（已排队的跳过）
func (e *Engine) Enqueue(tasks []Task) {
	e.mu.Lock()
	defer e.mu.Unlock()
	for _, t := range tasks {
		dup := false
		for _, q := range e.queue {
			if q.GID == t.GID && q.Token == t.Token {
				dup = true
				break
			}
		}
		if !dup {
			e.queue = append(e.queue, t)
		}
	}
}

// Start 开始跑队列（阻塞直到队列空或被取消）
func (e *Engine) Start(ctx context.Context) {
	e.mu.Lock()
	if e.running {
		e.mu.Unlock()
		return
	}
	e.running = true
	e.mu.Unlock()

	ctx, cancel := context.WithCancel(ctx)
	e.mu.Lock()
	e.cancel = cancel
	n := e.galleryWorkers
	e.mu.Unlock()

	defer func() {
		cancel()
		e.mu.Lock()
		e.running = false
		e.cancel = nil
		e.queue = nil
		e.mu.Unlock()
		if e.sink != nil {
			e.sink.Done()
		}
	}()

	// 并发 worker 共享一个 channel 队列
	work := make(chan Task)
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				case t, ok := <-work:
					if !ok {
						return
					}
					e.runTask(ctx, t)
				}
			}
		}()
	}

	// 派发
	go func() {
		defer close(work)
		for {
			e.mu.Lock()
			if len(e.queue) == 0 {
				e.mu.Unlock()
				return
			}
			t := e.queue[0]
			e.queue = e.queue[1:]
			e.mu.Unlock()

			select {
			case <-ctx.Done():
				return
			case work <- t:
			}
		}
	}()

	wg.Wait()
}

func (e *Engine) runTask(ctx context.Context, task Task) {
	task.Status = "running"
	task.Error = ""
	task.Done = 0
	task.Failed = 0
	task.Pages = map[string]string{}
	task.FailedList = []FailedItem{}
	e.save(task)

	// ① 画廊页（内部自动翻 ?p=N 分页收集全部页）
	gallery, err := e.client.FetchGalleryPage(ctx, task.GID, task.Token)
	if err != nil {
		if ctx.Err() != nil {
			e.markCanceled(task)
			return
		}
		e.fail(task, fmt.Errorf("解析画廊页失败: %w", err))
		return
	}

	task.Title = gallery.Title
	task.TotalPages = gallery.TotalPages
	task.Dir = BuildDir(e.store.Settings().DirTemplate, gallery)

	task.Status = "downloading"
	e.save(task)

	// 目录按模板生成
	root := e.store.Settings().DownloadRoot
	if root == "" {
		root = DefaultDownloadRoot()
	}
	dir := filepath.Join(root, task.Dir)

	// 老版本一律用「标题」当目录名。用户中途换了模板时，
	// 先看老的「标题目录」在不在 —— 在就继续往那儿下，
	// 否则换个模板就会把几千张图全部重下一遍。
	// ⚠️ 必须在 mkdir 之前判断，否则会白建出一个新目录。
	legacyDir := filepath.Join(root, SanitizeSegment(gallery.Title, task.GID))
	if dir != legacyDir {
		if _, err := os.Stat(legacyDir); err == nil {
			e.log("info", "沿用旧目录（模板改过，已有图在老位置）："+gallery.Title)
			dir = legacyDir
			rel, err := filepath.Rel(root, dir)
			if err == nil {
				task.Dir = rel
			}
		}
	}

	if err := os.MkdirAll(dir, 0o755); err != nil {
		e.fail(task, fmt.Errorf("创建目录失败: %w", err))
		return
	}
	e.save(task)

	// ② ③ 页级并发下载
	st := &pageState{
		task:   &task,
		dir:    dir,
		pages:  gallery.Pages,
		gid:    task.GID,
		token:  task.Token,
		title:  gallery.Title,
		total:  len(gallery.Pages),
		engine: e,
	}

	// 先做一遍快速扫描：已存在的直接算完成，剩下的才进队列
	st.scanExisting()
	st.saveManifest()

	// 页级 worker 池
	n := e.pageConcurrency
	if n > len(st.pending) {
		n = len(st.pending)
	}
	if n > 0 {
		var wg sync.WaitGroup
		work := make(chan eh.GalleryPageInfo, len(st.pending))
		for _, p := range st.pending {
			work <- p
		}
		close(work)

		for i := 0; i < n; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				for p := range work {
					st.downloadPage(ctx, p)
				}
			}()
		}
		wg.Wait()
	}

	// ---- 整轮扫尾：把刚失败的页再扫一遍 ----
	if st.failedCount() > 0 && ctx.Err() == nil {
		e.log("info", fmt.Sprintf("画廊 %s：%d 页首轮失败，开始扫尾重试", task.GID, st.failedCount()))

		retry := st.collectFailed()
		// 扫尾时清掉 failedList，让结果反映最终状态
		st.clearFailures()

		var wg sync.WaitGroup
		work := make(chan eh.GalleryPageInfo, len(retry))
		for _, p := range retry {
			work <- p
		}
		close(work)

		n := e.pageConcurrency
		if n > len(retry) {
			n = len(retry)
		}
		for i := 0; i < n; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				for p := range work {
					st.downloadPage(ctx, p)
				}
			}()
		}
		wg.Wait()

		if st.failedCount() > 0 {
			e.log("warn", fmt.Sprintf(
				"画廊 %s：扫尾后仍有 %d 页失败（换个线路再点「开始下载」可续传）",
				task.GID, st.failedCount()))
		} else {
			e.log("ok", fmt.Sprintf("画廊 %s：扫尾重试全部成功", task.GID))
		}
	}

	st.saveManifest()

	if ctx.Err() != nil {
		e.markCanceled(task)
		return
	}

	t := st.snapshot()
	if t.Failed > 0 {
		t.Status = "partial"
	} else {
		t.Status = "done"
	}
	e.save(t)
}

// pageState 是「一个画廊的页级下载状态」，所有字段都由 mutex 保护，
// 因为页级 worker 是并发的。
type pageState struct {
	mu     sync.Mutex
	engine *Engine

	task  *Task
	dir   string
	gid   string
	token string
	title string
	total int

	pages   []eh.GalleryPageInfo
	pending []eh.GalleryPageInfo

	done    map[int]bool
	failed  map[int]string // 页码 -> 错误
	failedL []FailedItem
}

// scanExisting 扫描已有文件（断点续传的第一依据：文件在且非空 = 完成）
func (s *pageState) scanExisting() {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.done = map[int]bool{}
	s.failed = map[int]string{}
	s.failedL = nil

	if s.task.Pages == nil {
		s.task.Pages = map[string]string{}
	}

	for _, p := range s.pages {
		fname := targetName(p, s.gid, s.total)
		dest := filepath.Join(s.dir, fname)
		if st, err := os.Stat(dest); err == nil && st.Size() > 0 {
			s.done[p.Page] = true
			s.task.Pages[strconv.Itoa(p.Page)] = "done"
		} else {
			s.pending = append(s.pending, p)
		}
	}
	s.task.Done = len(s.done)
}

// downloadPage 下单页（解析 URL → 下载），内部含 4 次退避重试
func (s *pageState) downloadPage(ctx context.Context, p eh.GalleryPageInfo) {
	if ctx.Err() != nil {
		return
	}

	fname := targetName(p, s.gid, s.total)
	dest := filepath.Join(s.dir, fname)

	err := s.engine.withRetry(ctx, func() error {
		imgURL, e := s.engine.client.ResolveImageURL(ctx, p.SPath)
		if e != nil {
			return e
		}
		_, e = s.engine.client.DownloadImage(ctx, imgURL, dest, nil)
		return e
	})

	s.mu.Lock()
	k := strconv.Itoa(p.Page)
	if err != nil {
		if ctx.Err() != nil || errors.Is(err, ErrCanceled) {
			s.mu.Unlock()
			s.markCanceled()
			return
		}
		s.failed[p.Page] = err.Error()
		s.failedL = append(s.failedL, FailedItem{Page: p.Page, File: fname, Error: err.Error()})
		s.task.Pages[k] = "failed"
	} else {
		s.done[p.Page] = true
		s.task.Pages[k] = "done"
	}
	s.syncCountsLocked()
	snap := s.snapshotLocked()
	s.mu.Unlock()

	s.engine.emitTask(snap)
}

func (s *pageState) syncCountsLocked() {
	s.task.Done = len(s.done)
	s.task.Failed = len(s.failed)
	s.task.FailedList = append([]FailedItem(nil), s.failedL...)
}

func (s *pageState) snapshotLocked() Task {
	s.syncCountsLocked()
	t := *s.task
	return t
}

func (s *pageState) snapshot() Task {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.snapshotLocked()
}

func (s *pageState) failedCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.failed)
}

func (s *pageState) clearFailures() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.failed = map[int]string{}
	s.failedL = nil
	s.task.Failed = 0
	s.task.FailedList = []FailedItem{}
}

// collectFailed 收集失败页（清空失败表前先取出）
func (s *pageState) collectFailed() []eh.GalleryPageInfo {
	s.mu.Lock()
	defer s.mu.Unlock()

	failed := intSet(keysOf(s.failed))
	var out []eh.GalleryPageInfo
	for _, p := range s.pages {
		if failed[p.Page] {
			out = append(out, p)
		}
	}
	return out
}

// markCanceled 任一页发现取消就整体标记
func (s *pageState) markCanceled() {
	s.engine.markCanceled(s.snapshot())
}

func (s *pageState) saveManifest() {
	s.mu.Lock()
	done := keysOf(s.done)
	failed := keysOf(s.failed)
	t := *s.task
	s.mu.Unlock()

	WriteManifest(s.dir, Manifest{
		GID:        s.gid,
		Token:      s.token,
		Title:      s.title,
		TotalPages: s.total,
		DonePages:  done,
		FailedPage: failed,
		CreatedAt:  t.CreatedAt,
	})
}

func keysOf[V any](m map[int]V) []int {
	out := make([]int, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	return out
}

// targetName 页码补零，保证排序
func targetName(p eh.GalleryPageInfo, gid string, total int) string {
	width := len(strconv.Itoa(total))
	n := strconv.Itoa(p.Page)
	for len(n) < width {
		n = "0" + n
	}
	return SanitizeFilename(p.Filename, gid+"-"+n+".webp")
}

// withRetry 指数退避重试
func (e *Engine) withRetry(ctx context.Context, fn func() error) error {
	const (
		tries = 4
		base  = 800 * time.Millisecond
		max   = 15 * time.Second
	)

	var lastErr error
	for i := 0; i < tries; i++ {
		if ctx.Err() != nil {
			return ErrCanceled
		}
		err := fn()
		if err == nil {
			return nil
		}
		if ctx.Err() != nil {
			return ErrCanceled
		}
		lastErr = err

		if i == tries-1 || !isRetryable(err) {
			return err
		}
		wait := time.Duration(float64(base) * math.Pow(2, float64(i)))
		if wait > max {
			wait = max
		}
		wait += time.Duration(rand.Intn(400)) * time.Millisecond
		e.log("warn", fmt.Sprintf("第 %d 次失败，%dms 后重试：%v", i+1, wait.Milliseconds(), err))

		select {
		case <-ctx.Done():
			return ErrCanceled
		case <-time.After(wait):
		}
	}
	return lastErr
}

// isRetryable 判断是否值得重试
func isRetryable(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, ErrCanceled) || errors.Is(err, context.Canceled) {
		return false
	}
	s := strings.ToLower(err.Error())
	for _, p := range retryablePatterns {
		if strings.Contains(s, p) {
			return true
		}
	}
	return false
}

// markCanceled 用户主动取消：退回排队，等续传
func (e *Engine) markCanceled(task Task) {
	task.Status = "queued"
	task.Error = ""
	e.save(task)
	e.log("warn", fmt.Sprintf("画廊 %s 已取消，再点开始可继续", task.GID))
}

func (e *Engine) fail(task Task, err error) {
	msg := err.Error()
	task.Status = "failed"
	task.Error = msg
	e.save(task)
	e.log("error", fmt.Sprintf("画廊 %s 失败：%s", task.GID, msg))
}

func (e *Engine) save(task Task) {
	task = e.store.UpsertTask(task)
	e.emitTask(task)
}

func (e *Engine) log(level, msg string) {
	if e.sink != nil {
		e.sink.Log(level, msg)
	}
}

// reSpaced 去掉路径里的多余空白（清洗用）
var reSpaced = regexp.MustCompile(`\s+`)

func cleanSpaces(s string) string { return reSpaced.ReplaceAllString(s, " ") }