hd4k-downloader/internal/core/store.go

package core

// 持久化：JSON 状态文件（原子写）
//
// 为什么不用 SQLite：断点续传的真正依据是「文件系统里文件在不在」，
// 状态文件只需记住画廊级别的进度。几百个任务的规模下 JSON 完全够用。
//
// 写文件用「临时文件 + rename」保证原子性。

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"hd4k-downloader/internal/hd4k"
)

// Site 站点配置（core 包里的镜像类型，避免前端到处引用 hd4k 包）
type Site = hd4k.Site

// Settings 用户设置
type Settings struct {
	DownloadRoot string `json:"downloadRoot"`
	Proxy        string `json:"proxy"`
	ProxyMode    string `json:"proxyMode"` // off | http | socks5
	// GalleryConcurrency 同时处理几个画廊
	GalleryConcurrency int `json:"galleryConcurrency"`
	// PageConcurrency 每个画廊内同时下几张图（总上限 = 两者相乘）
	PageConcurrency int `json:"pageConcurrency"`
	// RequestIntervalMs 全局请求间隔（毫秒，0 = 不限流）
	RequestIntervalMs int    `json:"requestIntervalMs"`
	Timeout           int    `json:"timeout"` // 毫秒
	BridgePort        int    `json:"bridgePort"`
	DirTemplate       string `json:"dirTemplate"`

	// Sites 支持的站点域名列表。换域名时在界面里加一条即可，不用改代码。
	Sites []Site `json:"sites"`

	// 内部字段：是否已探测过代理（不暴露给前端）
	Probed bool `json:"_probed"`
}

// FailedItem 单页失败记录
type FailedItem struct {
	Page  int    `json:"page"`
	File  string `json:"file"`
	Error string `json:"error"`
}

// Task 一个画廊的下载任务
//
// 唯一标识用 ID（域名 + 规范化路径），而不是 EH 那种 gid+token。
// 因为 4KHD 这类站点没有全局唯一 ID，URL 就是唯一标识；而且站点会换域名，
// 带上域名才能区分「同一个 slug 在不同域名上是两个画廊」。
type Task struct {
	ID     string `json:"id"`
	URL    string `json:"url"`
	Host   string `json:"host"`
	Path   string `json:"path"`
	Slug   string `json:"slug"`
	Raw    string `json:"raw"`
	Title  string `json:"title"`
	Status string `json:"status"` // queued|running|downloading|partial|done|failed
	Done   int    `json:"done"`
	Failed int    `json:"failed"`
	// Total 图片总数（不是 HTML 分页数）
	Total      int               `json:"total"`
	Error      string            `json:"error"`
	Dir        string            `json:"dir"`
	Pages      map[string]string `json:"pages"`
	FailedList []FailedItem      `json:"failedList"`
	Selected   bool              `json:"selected"`
	CreatedAt  int64             `json:"createdAt"`
	UpdatedAt  int64             `json:"updatedAt"`
}

// State 状态文件顶层结构
type State struct {
	Settings Settings `json:"settings"`
	Tasks    []Task   `json:"tasks"`
}

// DefaultSettings 返回默认设置
func DefaultSettings() Settings {
	return Settings{
		DownloadRoot:       "",
		Proxy:              "",
		ProxyMode:          "off",
		GalleryConcurrency: 5,
		PageConcurrency:    8,
		RequestIntervalMs:  0,
		Timeout:            30000,
		BridgePort:         37778,
		DirTemplate:        "{title}",
		Sites:              hd4k.DefaultSites(),
	}
}

// Store 是带锁的 JSON 状态存储，并发安全
type Store struct {
	mu    sync.RWMutex
	path  string
	state State

	flushTimer *time.Timer
	dirty      bool
}

// NewStore 打开（或新建）状态文件
func NewStore(filePath string) *Store {
	s := &Store{path: filePath}
	s.load()
	return s
}

func (s *Store) load() {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.state = State{Settings: DefaultSettings(), Tasks: []Task{}}

	raw, err := os.ReadFile(s.path)
	if err != nil {
		return
	}
	var loaded State
	if json.Unmarshal(raw, &loaded) != nil {
		return
	}

	// 补齐缺失字段（向后兼容 / 手工改坏配置时不至于崩）
	def := DefaultSettings()
	if loaded.Settings.DownloadRoot == "" {
		loaded.Settings.DownloadRoot = def.DownloadRoot
	}
	if loaded.Settings.GalleryConcurrency <= 0 {
		loaded.Settings.GalleryConcurrency = def.GalleryConcurrency
	}
	if loaded.Settings.PageConcurrency <= 0 {
		loaded.Settings.PageConcurrency = def.PageConcurrency
	}
	if loaded.Settings.RequestIntervalMs < 0 {
		loaded.Settings.RequestIntervalMs = 0
	}
	if loaded.Settings.Timeout <= 0 {
		loaded.Settings.Timeout = def.Timeout
	}
	if loaded.Settings.BridgePort <= 0 {
		loaded.Settings.BridgePort = def.BridgePort
	}
	if loaded.Settings.DirTemplate == "" {
		loaded.Settings.DirTemplate = def.DirTemplate
	}
	if len(loaded.Settings.Sites) == 0 {
		loaded.Settings.Sites = def.Sites
	} else {
		// 迁移：旧版 state.json 里存的站点没有 aliases 字段，
		// 而别名归一（www.4khd.com → migmg.ssuu.uk）靠的就是它。
		// 不补的话升级后老用户永远用不上别名 —— 实测同一个画廊会因为
		// 「从哪个域名复制的链接」变成两个不同任务。
		// 只补 aliases/scheme，用户的 label/enabled 保持不动。
		loaded.Settings.Sites = mergeSiteAliases(loaded.Settings.Sites, def.Sites)
	}
	if loaded.Tasks == nil {
		loaded.Tasks = []Task{}
	}
	s.state = loaded
}

// mergeSiteAliases 把默认站点里的 aliases/scheme 补进已持久化的站点列表。
//
// 老版本的 state.json 没有 aliases 字段。别名归一（www.4khd.com → migmg.ssuu.uk）
// 全靠它 —— 不补的话升级后同一个画廊会因为「从哪个域名复制的链接」变成两个任务。
// 用户自己的 label / enabled / scheme 一律保留，只在为空时才用默认值。
func mergeSiteAliases(have, def []Site) []Site {
	byHost := map[string]int{}
	for i, s := range have {
		byHost[strings.ToLower(strings.TrimSpace(s.Host))] = i
	}
	for _, d := range def {
		key := strings.ToLower(strings.TrimSpace(d.Host))
		i, ok := byHost[key]
		if !ok {
			// 默认站点不在用户列表里 —— 用户可能删过，尊重用户的选择，不擅自加回
			continue
		}
		if len(have[i].Aliases) == 0 && len(d.Aliases) > 0 {
			have[i].Aliases = append([]string(nil), d.Aliases...)
		}
		if have[i].Scheme == "" && d.Scheme != "" {
			have[i].Scheme = d.Scheme
		}
		if have[i].Label == "" {
			have[i].Label = d.Label
		}
	}
	return have
}

// Flush 原子写：先写临时文件再 rename
func (s *Store) Flush() {
	s.mu.RLock()
	data, err := json.MarshalIndent(s.state, "", "  ")
	path := s.path
	s.dirty = false
	s.mu.RUnlock()

	if err != nil {
		return
	}
	tmp := path + ".tmp"
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return
	}
	if err := os.WriteFile(tmp, data, 0o644); err != nil {
		return
	}
	_ = os.Rename(tmp, path)
}

// ScheduleFlush 合并短时间内的多次写，降低磁盘压力
func (s *Store) ScheduleFlush(delay int) {
	if delay <= 0 {
		delay = 400
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.flushTimer != nil {
		return
	}
	s.dirty = true
	s.flushTimer = time.AfterFunc(time.Duration(delay)*time.Millisecond, func() {
		s.mu.Lock()
		s.flushTimer = nil
		s.mu.Unlock()
		s.Flush()
	})
}

// Stop 停掉待执行的延迟写（退出前调用）
func (s *Store) Stop() {
	s.mu.Lock()
	if s.flushTimer != nil {
		s.flushTimer.Stop()
		s.flushTimer = nil
	}
	s.mu.Unlock()
}

// Settings 取当前设置（返回副本）
func (s *Store) Settings() Settings {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := s.state.Settings
	out.Sites = append([]Site(nil), s.state.Settings.Sites...)
	return out
}

// UpdateSettings 合并设置
func (s *Store) UpdateSettings(patch map[string]interface{}) Settings {
	s.mu.Lock()
	// 注意：这里不能用 defer Unlock。返回值必须在解锁**之后**取，
	// 而 Go 的 RWMutex 不可重入 —— 持写锁时调 Settings() 会再 RLock
	// 同一把锁，永久死锁。症状是「点保存设置后整个程序卡死且不报错」。
	// 所以改成手动解锁。

	// 站点列表是结构体，切出来单独处理，避免走 map[string]interface{}
	// 那条路时被 JSON 解成 []any 导致反序列化失败
	if rawSites, ok := patch["sites"]; ok {
		delete(patch, "sites")
		if b, err := json.Marshal(rawSites); err == nil {
			var sites []Site
			if json.Unmarshal(b, &sites) == nil {
				s.state.Settings.Sites = sites
			}
		}
	}

	m, _ := json.Marshal(s.state.Settings)
	var cur map[string]interface{}
	_ = json.Unmarshal(m, &cur)
	for k, v := range patch {
		cur[k] = v
	}
	m2, _ := json.Marshal(cur)
	var next Settings
	if json.Unmarshal(m2, &next) == nil {
		if len(next.Sites) == 0 {
			next.Sites = s.state.Settings.Sites
		}
		s.state.Settings = next
	}
	s.startFlushLocked()

	// 取返回值：拷一份 Sites 切片，别把内部切片漏给调用方
	// （调用方改它就等于绕过锁直接改内部状态）。
	out := s.state.Settings
	out.Sites = append([]Site(nil), s.state.Settings.Sites...)
	s.mu.Unlock()
	return out
}

func (s *Store) startFlushLocked() {
	if s.flushTimer != nil {
		return
	}
	s.dirty = true
	s.flushTimer = time.AfterFunc(400*time.Millisecond, func() {
		s.mu.Lock()
		s.flushTimer = nil
		s.mu.Unlock()
		s.Flush()
	})
}

// Tasks 取全部任务（深拷贝，防止调用方改到内部状态）
func (s *Store) Tasks() []Task {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make([]Task, 0, len(s.state.Tasks))
	for _, t := range s.state.Tasks {
		out = append(out, t.clone())
	}
	return out
}

func (t Task) clone() Task {
	c := t
	if t.Pages != nil {
		c.Pages = make(map[string]string, len(t.Pages))
		for k, v := range t.Pages {
			c.Pages[k] = v
		}
	}
	if t.FailedList != nil {
		c.FailedList = append([]FailedItem(nil), t.FailedList...)
	}
	return c
}

// FindTask 按 ID 找任务
func (s *Store) FindTask(id string) *Task {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for i := range s.state.Tasks {
		if s.state.Tasks[i].ID == id {
			c := s.state.Tasks[i].clone()
			return &c
		}
	}
	return nil
}

// AddTasks 批量加任务（已存在的跳过），返回新增数
func (s *Store) AddTasks(refs []hd4k.GalleryRef) int {
	s.mu.Lock()
	defer s.mu.Unlock()

	added := 0
	now := time.Now().UnixMilli()
	for _, r := range refs {
		if r.ID == "" {
			continue
		}
		exists := false
		for _, t := range s.state.Tasks {
			if t.ID == r.ID {
				exists = true
				break
			}
		}
		if exists {
			continue
		}
		s.state.Tasks = append([]Task{{
			ID:         r.ID,
			URL:        r.URL,
			Host:       r.Host,
			Path:       r.Path,
			Slug:       r.Slug,
			Raw:        r.Raw,
			Title:      r.Title,
			Status:     "queued",
			Pages:      map[string]string{},
			FailedList: []FailedItem{},
			CreatedAt:  now,
			UpdatedAt:  now,
		}}, s.state.Tasks...)
		added++
	}
	s.startFlushLocked()
	return added
}

// SetTitle 给已有任务补标题（油猴推送时带标题来的）
func (s *Store) SetTitle(id, title string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := range s.state.Tasks {
		if s.state.Tasks[i].ID == id {
			if s.state.Tasks[i].Title == "" && title != "" {
				s.state.Tasks[i].Title = truncate(title, 300)
				s.startFlushLocked()
				return true
			}
			return false
		}
	}
	return false
}

func truncate(s string, n int) string {
	r := []rune(s)
	if len(r) <= n {
		return s
	}
	return string(r[:n])
}

// UpsertTask 更新/插入任务。
//
// ⚠️ 存进去前必须 clone() —— 否则外部（图片级 worker）继续改同一个 map，
// 而 store 的 Flush() 会在后台 json.Marshal 整个 state —— 撞上就是
// fatal error: concurrent map write（不可 recover，进程闪退）。
func (s *Store) UpsertTask(t Task) Task {
	s.mu.Lock()
	defer s.mu.Unlock()

	t.UpdatedAt = time.Now().UnixMilli()
	if t.Pages == nil {
		t.Pages = map[string]string{}
	}
	if t.FailedList == nil {
		t.FailedList = []FailedItem{}
	}

	for i := range s.state.Tasks {
		if s.state.Tasks[i].ID == t.ID {
			if t.CreatedAt == 0 {
				t.CreatedAt = s.state.Tasks[i].CreatedAt
			}
			s.state.Tasks[i] = t.clone()
			s.startFlushLocked()
			return t.clone()
		}
	}

	if t.CreatedAt == 0 {
		t.CreatedAt = t.UpdatedAt
	}
	stored := t.clone()
	s.state.Tasks = append([]Task{stored}, s.state.Tasks...)
	s.startFlushLocked()
	return stored.clone()
}

// RemoveTask 删单个任务
func (s *Store) RemoveTask(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := range s.state.Tasks {
		if s.state.Tasks[i].ID == id {
			s.state.Tasks = append(s.state.Tasks[:i], s.state.Tasks[i+1:]...)
			s.startFlushLocked()
			return
		}
	}
}

// ClearAll 清空任务列表（或只删勾选的）。
// ⚠️ 只清任务列表，本地已下载的文件一个都不动。
func (s *Store) ClearAll(onlySelected bool) int {
	s.mu.Lock()
	defer s.mu.Unlock()

	before := len(s.state.Tasks)
	kept := make([]Task, 0, before)
	for _, t := range s.state.Tasks {
		if onlySelected && !t.Selected {
			kept = append(kept, t)
		}
	}
	s.state.Tasks = kept
	s.startFlushLocked()
	return before - len(s.state.Tasks)
}

// ClearFinished 只清 done。partial/failed 保留。
func (s *Store) ClearFinished() int {
	s.mu.Lock()
	defer s.mu.Unlock()

	before := len(s.state.Tasks)
	kept := make([]Task, 0, before)
	for _, t := range s.state.Tasks {
		if t.Status != "done" {
			kept = append(kept, t)
		}
	}
	s.state.Tasks = kept
	s.startFlushLocked()
	return before - len(s.state.Tasks)
}

// CountDone 数一下已完成
func (s *Store) CountDone() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	n := 0
	for _, t := range s.state.Tasks {
		if t.Status == "done" {
			n++
		}
	}
	return n
}

// RequeueFailed 把失败/部分失败/中断的重新排队
func (s *Store) RequeueFailed() int {
	s.mu.Lock()
	defer s.mu.Unlock()

	n := 0
	for i := range s.state.Tasks {
		switch s.state.Tasks[i].Status {
		case "failed", "partial", "downloading", "running", "parsing":
			s.state.Tasks[i].Status = "queued"
			s.state.Tasks[i].Error = ""
			n++
		}
	}
	if n > 0 {
		s.startFlushLocked()
	}
	return n
}

// RecoverOnStart 启动时把残留的 running/downloading/parsing 归位
func (s *Store) RecoverOnStart() int {
	s.mu.Lock()
	defer s.mu.Unlock()

	n := 0
	for i := range s.state.Tasks {
		switch s.state.Tasks[i].Status {
		case "running", "downloading", "parsing":
			s.state.Tasks[i].Status = "queued"
			n++
		}
	}
	if n > 0 {
		s.startFlushLocked()
	}
	return n
}

// SelectAll 标记勾选状态
func (s *Store) SelectAll(ids []string) []Task {
	set := make(map[string]bool, len(ids))
	for _, g := range ids {
		set[g] = true
	}
	s.mu.Lock()
	for i := range s.state.Tasks {
		s.state.Tasks[i].Selected = set[s.state.Tasks[i].ID]
	}
	s.startFlushLocked()
	out := make([]Task, 0, len(s.state.Tasks))
	for _, t := range s.state.Tasks {
		out = append(out, t.clone())
	}
	s.mu.Unlock()
	return out
}