app.go

package main

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	goruntime "runtime"
	"strings"
	"sync"

	wruntime "github.com/wailsapp/wails/v2/pkg/runtime"

	"kstudio/internal/config"
	"kstudio/internal/engine"
	"kstudio/internal/imagegen"
	"kstudio/internal/llm"
	"kstudio/internal/store"
)

// App is the API surface exposed to the frontend.
type App struct {
	ctx context.Context

	mu     sync.RWMutex
	cfg    config.Config
	store  *store.Store
	engine *engine.Engine
}

// NewApp wires the app together from the persisted settings.
func NewApp() *App {
	// 软件改名前的设置目录还叫 QwenImageStudio，先搬过来再读，免得老用户开机发现
	// Key 和作品目录都空了。已经搬过就什么都不做。
	config.Migrate()
	cfg := config.Load()
	a := &App{cfg: cfg, store: store.New(cfg.OutputDir)}
	a.engine = engine.New(a.store, a.currentConfig, a.emit)
	return a
}

func (a *App) startup(ctx context.Context) {
	a.ctx = ctx
	if root := a.store.Root(); root != "" {
		if err := os.MkdirAll(root, 0o755); err != nil {
			fmt.Println("无法创建输出目录:", err)
		}
	}
}

func (a *App) currentConfig() config.Config {
	a.mu.RLock()
	defer a.mu.RUnlock()
	return a.cfg
}

func (a *App) emit(event string, data any) {
	if a.ctx == nil {
		return
	}
	wruntime.EventsEmit(a.ctx, event, data)
}

// ---------------------------------------------------------------- settings

// TestResult is the outcome of a connection test.
type TestResult struct {
	OK      bool   `json:"ok"`
	Message string `json:"message"`
	URL     string `json:"url"`
}

// GetConfig returns the current settings.
func (a *App) GetConfig() config.Config { return a.currentConfig() }

// SaveConfig persists settings and applies them immediately.
func (a *App) SaveConfig(cfg config.Config) (config.Config, error) {
	normalized := cfg.Normalized()
	if err := config.Save(normalized); err != nil {
		return normalized, err
	}
	a.mu.Lock()
	a.cfg = normalized
	a.mu.Unlock()
	a.store.SetRoot(normalized.OutputDir)
	if err := os.MkdirAll(normalized.OutputDir, 0o755); err != nil {
		return normalized, fmt.Errorf("输出目录不可用：%w", err)
	}
	return normalized, nil
}

// ResolveEndpoints shows the exact URLs the two providers will be called on,
// so the user can see what a base URL expands to before saving.
func (a *App) ResolveEndpoints(cfg config.Config) map[string]string {
	return map[string]string{
		"orchestrator": config.ChatCompletionsURL(cfg.Orchestrator.BaseURL),
		"image":        config.ImageURL(cfg.Image),
		"imageProto":   imagegen.ProtocolLabel(config.ImageSpec(cfg.Image).Protocol),
	}
}

// TestOrchestrator sends the cheapest possible completion to the decision model.
func (a *App) TestOrchestrator(p config.Provider) TestResult {
	target := config.ChatCompletionsURL(p.BaseURL)
	p.APIKey = config.CleanKey(p.APIKey)
	if problem := config.KeyProblem(p.APIKey); problem != "" {
		return TestResult{Message: problem, URL: target}
	}
	client := llm.New(target, p.APIKey, p.Model, 64, 0)
	msg, err := client.Ping(a.context())
	if err != nil {
		return TestResult{Message: withKeyHint(err.Error(), p.APIKey), URL: target}
	}
	return TestResult{OK: true, Message: msg, URL: target}
}

// withKeyHint appends what we actually sent when the endpoint rejects the
// credentials, so a swapped or truncated key is obvious from the message.
func withKeyHint(msg, key string) string {
	low := strings.ToLower(msg)
	auth := strings.Contains(msg, "401") || strings.Contains(msg, "403") ||
		strings.Contains(low, "authentication") || strings.Contains(low, "invalid api") ||
		strings.Contains(low, "invalidapikey") || strings.Contains(low, "unauthorized") ||
		strings.Contains(msg, "鉴权")
	if !auth {
		return msg
	}
	return msg + "\n本次发送的 Key：" + config.MaskKey(key) + "。请确认这一栏填的是 API Key 本身，而且和 Base URL 属于同一家服务商。"
}

// TestImage validates the image endpoint without spending money: it sends an
// intentionally incomplete body, so a parameter complaint proves auth is fine.
// A custom template is the one exception — sending it would be a real, billed
// generation, so it is only validated offline.
func (a *App) TestImage(p config.ImageProvider) TestResult {
	spec := config.ImageSpec(p)
	if spec.Protocol != imagegen.ProtoCustom || spec.APIKey != "" {
		if problem := config.KeyProblem(spec.APIKey); problem != "" {
			return TestResult{Message: problem, URL: spec.URL}
		}
	}
	msg, err := imagegen.New(spec).Ping(a.context())
	if err != nil {
		return TestResult{Message: withKeyHint(err.Error(), spec.APIKey), URL: spec.URL}
	}
	return TestResult{OK: true, Message: msg, URL: spec.URL}
}

// PickOutputDir opens a folder picker for the output directory.
func (a *App) PickOutputDir() (string, error) {
	dir, err := wruntime.OpenDirectoryDialog(a.ctx, wruntime.OpenDialogOptions{
		Title:                "选择作品输出目录",
		DefaultDirectory:     a.store.Root(),
		CanCreateDirectories: true,
	})
	if err != nil {
		return "", err
	}
	return dir, nil
}

// AppInfo is small metadata for the settings screen.
func (a *App) AppInfo() map[string]string {
	return map[string]string{
		"version":    appVersion,
		"configPath": config.Path(),
		"outputDir":  a.store.Root(),
		"go":         goruntime.Version(),
		"ratios":     strings.Join(imagegen.Ratios(), ","),
		"tiers":      strings.Join(imagegen.Tiers(), ","),
		"sizeRange":  fmt.Sprintf("%d-%d", imagegen.SizeMin, imagegen.SizeMax),
	}
}

// SizeTable maps every resolution tier and ratio to the pixel size that will
// be requested, so the UI can label a choice without a call per keystroke.
func (a *App) SizeTable() map[string]map[string]string { return imagegen.SizeTable() }

// ProtocolOption is one selectable way of reaching an image provider.
type ProtocolOption struct {
	Value string `json:"value"`
	Label string `json:"label"`
	Hint  string `json:"hint"`
	// Refs is false when this protocol has nowhere to carry a reference image,
	// so the UI can say that the first-image style lock will not apply. Every
	// preset can carry one now; only a custom template may leave the slot out,
	// and the UI decides that from the template itself.
	Refs bool `json:"refs"`
	// Custom marks the one option that needs the request template filled in.
	Custom bool `json:"custom"`
}

// protocolHints explains each option in the words of the thing the user is
// looking at — a provider's docs page — rather than in ours.
var protocolHints = map[string]string{
	imagegen.ProtoAuto:          "按 Base URL 自动判断：填百炼地址就走 qwen-image，填中转站地址就走 OpenAI 图像接口",
	imagegen.ProtoDashScope:     "阿里云百炼 qwen-image-3.0(-pro)：multimodal-generation，直接返回图片链接",
	imagegen.ProtoDashScopeTask: "百炼万相 wanx / wan2.2 等异步文生图：先提交任务，再自动轮询结果",
	imagegen.ProtoOpenAIImages:  "POST /v1/images/generations：OpenAI、SiliconFlow、ModelScope、各类中转站；带参考图时自动改发 /v1/images/edits（模型名要写 gpt-image-2 这类接口自己的名字）",
	imagegen.ProtoOpenAIChat:    "把出图模型挂在 /v1/chat/completions 上：多数中转站的 gemini-image、gpt-image 是这种",
	imagegen.ProtoGemini:        "Google 官方 generateContent：图片以 base64 内联返回，参考图会自动转成内联数据",
	imagegen.ProtoCustom:        "上面都不匹配时，自己写请求体模板，用 {{prompt}} 之类的占位符引用参数",
}

// ImageProtocols lists every way the app can reach an image provider, so the
// settings screen can offer them without hard-coding the list in JavaScript.
func (a *App) ImageProtocols() []ProtocolOption {
	all := imagegen.Protocols()
	out := make([]ProtocolOption, 0, len(all))
	for _, p := range all {
		label := imagegen.ProtocolLabel(p)
		if p == imagegen.ProtoAuto {
			label = "自动识别"
		}
		out = append(out, ProtocolOption{
			Value:  p,
			Label:  label,
			Hint:   protocolHints[p],
			Refs:   imagegen.ProtocolRefs(p),
			Custom: p == imagegen.ProtoCustom,
		})
	}
	return out
}

// DefaultCustomTemplate is the starting request body for a custom provider, so
// the UI's "恢复默认模板" button does not have to carry a copy of it.
func (a *App) DefaultCustomTemplate() string { return imagegen.DefaultCustomBody }

// OpenConfigFolder reveals config.json in the file manager.
func (a *App) OpenConfigFolder() error { return revealPath(config.Path()) }

// ---------------------------------------------------------------- projects

// StartGeneration launches a run and returns the freshly created project.
func (a *App) StartGeneration(req engine.StartRequest) (*store.Project, error) {
	return a.engine.Start(req)
}

// CancelRun stops an in-flight run.
func (a *App) CancelRun(id string) error { return a.engine.Cancel(id) }

// IsRunning reports whether a project has work in flight.
func (a *App) IsRunning(id string) bool { return a.engine.Running(id) }

// RegenerateShot re-renders one image, optionally with an edited prompt.
func (a *App) RegenerateShot(id string, index int, prompt string) error {
	return a.engine.RegenerateShot(id, index, prompt)
}

// RetryFailedShots re-renders every image of a project that is not done yet and
// returns how many were queued.
func (a *App) RetryFailedShots(id string) (int, error) { return a.engine.RetryFailed(id) }

// ListProjects returns all projects, newest first.
func (a *App) ListProjects() ([]store.Summary, error) { return a.store.List() }

// GetProject loads one project, repairing state left behind by a crash.
func (a *App) GetProject(id string) (*store.Project, error) {
	p, err := a.store.Load(id)
	if err != nil {
		return nil, err
	}
	a.engine.Reconcile(p)
	return p, nil
}

// DeleteProject removes a project folder from disk.
func (a *App) DeleteProject(id string) error {
	if a.engine.Running(id) {
		return errors.New("该项目正在生成中，请先取消任务")
	}
	return a.store.Delete(id)
}

// RenameProject changes a project's display title.
func (a *App) RenameProject(id, title string) (*store.Project, error) {
	return a.store.Rename(id, title)
}

// Gallery lists every finished image across all projects.
func (a *App) Gallery() ([]store.GalleryItem, error) { return a.store.Gallery() }

// OpenProjectFolder opens the project folder in the file manager.
func (a *App) OpenProjectFolder(id string) error {
	dir, err := a.store.DirOf(id)
	if err != nil {
		return err
	}
	return openPath(dir)
}

// RevealImage highlights one generated image in the file manager.
func (a *App) RevealImage(id, file string) error {
	path, err := a.store.ImagePath(id, file)
	if err != nil {
		return err
	}
	return revealPath(path)
}

// ExportProject copies the finished images plus the markdown shot list into a
// folder the user picks.
func (a *App) ExportProject(id string) (string, error) {
	dest, err := wruntime.OpenDirectoryDialog(a.ctx, wruntime.OpenDialogOptions{
		Title:                "选择导出位置",
		CanCreateDirectories: true,
	})
	if err != nil {
		return "", err
	}
	if dest == "" {
		return "", nil // user cancelled
	}
	out, count, err := a.store.ExportTo(id, dest)
	if err != nil {
		return "", err
	}
	if count == 0 {
		return out, errors.New("该项目还没有可导出的图片，仅导出了分镜说明")
	}
	return out, nil
}

// ExportImage saves one generated image wherever the user points. The default
// filename says which project and which shot it came from, so the picture is
// still identifiable once it's out of the app.
func (a *App) ExportImage(id, file string) (string, error) {
	if _, err := a.store.ImagePath(id, file); err != nil {
		return "", err
	}
	dest, err := wruntime.SaveFileDialog(a.ctx, wruntime.SaveDialogOptions{
		Title:                "导出这张图片",
		DefaultFilename:      a.store.ExportName(id, file),
		CanCreateDirectories: true,
		Filters: []wruntime.FileFilter{
			{DisplayName: "PNG 图片 (*.png)", Pattern: "*.png"},
			{DisplayName: "所有文件 (*.*)", Pattern: "*.*"},
		},
	})
	if err != nil {
		return "", err
	}
	if dest == "" {
		return "", nil // user cancelled
	}
	if err := a.store.ExportImageTo(id, file, dest); err != nil {
		return "", err
	}
	return dest, nil
}

// CopyText puts text on the system clipboard.
func (a *App) CopyText(text string) error {
	return wruntime.ClipboardSetText(a.ctx, text)
}

// OpenURL opens a link in the default browser.
func (a *App) OpenURL(target string) error {
	wruntime.BrowserOpenURL(a.ctx, target)
	return nil
}

func (a *App) context() context.Context {
	if a.ctx != nil {
		return a.ctx
	}
	return context.Background()
}

// ---------------------------------------------------------------- media

// mediaHandler serves generated images to the webview from disk under
// /media/<projectID>/<file>. Anything outside the output directory is refused.
func (a *App) mediaHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		rest := strings.TrimPrefix(r.URL.Path, "/media/")
		if rest == r.URL.Path {
			http.NotFound(w, r)
			return
		}
		parts := strings.SplitN(strings.Trim(rest, "/"), "/", 2)
		if len(parts) != 2 {
			http.Error(w, "bad media path", http.StatusBadRequest)
			return
		}
		id, err := url.PathUnescape(parts[0])
		if err != nil {
			http.Error(w, "bad project id", http.StatusBadRequest)
			return
		}
		file, err := url.PathUnescape(parts[1])
		if err != nil {
			http.Error(w, "bad file name", http.StatusBadRequest)
			return
		}
		path, err := a.store.ImagePath(id, file)
		if err != nil {
			http.Error(w, err.Error(), http.StatusForbidden)
			return
		}
		info, err := os.Stat(path)
		if err != nil || info.IsDir() {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Cache-Control", "no-cache")
		http.ServeFile(w, r, path)
	})
}

// ---------------------------------------------------------------- os helpers

func openPath(path string) error {
	path = filepath.Clean(path)
	switch goruntime.GOOS {
	case "windows":
		// explorer.exe exits with status 1 even when it succeeds.
		_ = exec.Command("explorer", path).Run()
		return nil
	case "darwin":
		return exec.Command("open", path).Run()
	default:
		return exec.Command("xdg-open", path).Run()
	}
}

func revealPath(path string) error {
	path = filepath.Clean(path)
	switch goruntime.GOOS {
	case "windows":
		_ = exec.Command("explorer", "/select,"+path).Run()
		return nil
	case "darwin":
		return exec.Command("open", "-R", path).Run()
	default:
		return openPath(filepath.Dir(path))
	}
}