dingtalk-workspace-cli/internal/helpers/import_flow.go

// Copyright 2026 Alibaba Group
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package helpers

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/spf13/cobra"
)

const importMaxFileSize int64 = 20 * 1024 * 1024

type importPollTimeoutError struct {
	taskID   string
	maxPolls int
}

func (e *importPollTimeoutError) Error() string {
	return fmt.Sprintf("导入任务超时：已轮询 %d 次仍在处理中 (taskId=%s)", e.maxPolls, e.taskID)
}

type importPollPolicy struct {
	maxPolls int
	interval func(attempt int) time.Duration
	wait     func(context.Context, time.Duration) error
}

type importFlowConfig struct {
	operation            string
	queryOperation       string
	supportedFormats     map[string]bool
	supportedFormatsText string
	folderFlags          []string
	workspaceFlags       []string
	requireTarget        bool
	serverID             string
	includeNodeID        bool
	timeoutAsResult      bool
	nextCommand          string
	poll                 importPollPolicy
}

type preparedImportFile struct {
	path      string
	name      string
	extension string
	size      int64
	folder    string
	workspace string
}

func defaultImportPollPolicy() importPollPolicy {
	return importPollPolicy{
		maxPolls: 30,
		interval: func(attempt int) time.Duration {
			switch {
			case attempt <= 5:
				return 2 * time.Second
			case attempt <= 10:
				return 5 * time.Second
			case attempt <= 20:
				return 10 * time.Second
			default:
				return 15 * time.Second
			}
		},
		wait: waitForImportPoll,
	}
}

func docImportFlowConfig() importFlowConfig {
	return importFlowConfig{
		operation:            "导入本地文件为在线文档",
		queryOperation:       "查询导入任务结果",
		supportedFormats:     map[string]bool{"docx": true, "doc": true, "xlsx": true, "xls": true, "md": true, "txt": true, "xmind": true, "mark": true},
		supportedFormatsText: "docx, doc, xlsx, xls, md, txt, xmind, mark",
		folderFlags:          []string{"folder", "folder-id"},
		workspaceFlags:       []string{"workspace", "workspace-id"},
		nextCommand:          "dws doc import get --task-id %s",
		poll:                 defaultImportPollPolicy(),
	}
}

func sheetImportFlowConfig() importFlowConfig {
	return importFlowConfig{
		operation:            "导入本地表格文件为在线电子表格",
		queryOperation:       "查询表格导入任务结果",
		supportedFormats:     map[string]bool{"xlsx": true, "xls": true},
		supportedFormatsText: "xlsx, xls",
		folderFlags:          []string{"folder-token", "folder"},
		workspaceFlags:       []string{"workspace"},
		requireTarget:        true,
		serverID:             "doc",
		includeNodeID:        true,
		timeoutAsResult:      true,
		nextCommand:          "dws sheet import get --task-id %s",
		poll:                 defaultImportPollPolicy(),
	}
}

func waitForImportPoll(ctx context.Context, duration time.Duration) error {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		return nil
	}
}

func importFlagValue(cmd *cobra.Command, names ...string) string {
	for _, name := range names {
		if cmd.Flags().Lookup(name) == nil {
			continue
		}
		if value, _ := cmd.Flags().GetString(name); value != "" {
			return value
		}
	}
	return ""
}

func prepareImportFile(cmd *cobra.Command, args []string, cfg importFlowConfig) (preparedImportFile, error) {
	filePath := mustGetFlag(cmd, "file")
	if filePath == "" && len(args) > 0 {
		filePath = args[0]
	}
	if filePath == "" {
		return preparedImportFile{}, fmt.Errorf("flag --file is required (or pass file path as argument)")
	}

	fileInfo, err := os.Stat(filePath)
	if err != nil {
		return preparedImportFile{}, fmt.Errorf("cannot read file %s: %w", filePath, err)
	}
	if fileInfo.IsDir() {
		return preparedImportFile{}, fmt.Errorf("%s is a directory, not a file", filePath)
	}
	if fileInfo.Size() > importMaxFileSize {
		return preparedImportFile{}, fmt.Errorf("file size %d bytes exceeds 20MB limit", fileInfo.Size())
	}
	if fileInfo.Size() == 0 {
		return preparedImportFile{}, fmt.Errorf("file is empty: %s", filePath)
	}

	extension := strings.TrimPrefix(strings.ToLower(filepath.Ext(filePath)), ".")
	if !cfg.supportedFormats[extension] {
		return preparedImportFile{}, fmt.Errorf("unsupported file format %q, supported: %s", extension, cfg.supportedFormatsText)
	}

	name, _ := cmd.Flags().GetString("name")
	if name == "" {
		fileName := filepath.Base(filePath)
		name = strings.TrimSuffix(fileName, filepath.Ext(fileName))
	}
	folder := importFlagValue(cmd, cfg.folderFlags...)
	workspace := importFlagValue(cmd, cfg.workspaceFlags...)
	if cfg.requireTarget && folder == "" && workspace == "" {
		return preparedImportFile{}, fmt.Errorf("--folder-token 与 --workspace 至少需要提供一个（导入目标位置）")
	}

	return preparedImportFile{
		path:      filePath,
		name:      name,
		extension: extension,
		size:      fileInfo.Size(),
		folder:    folder,
		workspace: workspace,
	}, nil
}

func (cfg importFlowConfig) callTool(ctx context.Context, toolName string, args map[string]any) (string, error) {
	if cfg.serverID != "" {
		return callMCPToolReturnTextOnServer(ctx, cfg.serverID, toolName, args)
	}
	return callMCPToolReturnText(ctx, toolName, args)
}

func runImportCommand(cmd *cobra.Command, args []string, cfg importFlowConfig) error {
	file, err := prepareImportFile(cmd, args, cfg)
	if err != nil {
		return err
	}

	if deps.Caller.DryRun() {
		deps.Out.PrintKeyValue("操作", cfg.operation)
		deps.Out.PrintKeyValue("文件", file.path)
		deps.Out.PrintKeyValue("名称", file.name)
		deps.Out.PrintKeyValue("格式", file.extension)
		deps.Out.PrintKeyValue("大小", fmt.Sprintf("%d bytes", file.size))
		return nil
	}

	ctx := cmd.Context()
	if ctx == nil {
		ctx = context.Background()
	}

	deps.Out.PrintInfo("[1/4] 创建导入会话...")
	sessionArgs := map[string]any{
		"fileName": file.name,
		"suffix":   file.extension,
		"fileSize": file.size,
	}
	if file.folder != "" {
		sessionArgs["targetFolderId"] = file.folder
	}
	if file.workspace != "" {
		sessionArgs["workspaceId"] = file.workspace
	}

	sessionText, err := cfg.callTool(ctx, "create_import_session", sessionArgs)
	if err != nil {
		return fmt.Errorf("创建导入会话失败: %w", err)
	}
	var sessionResult map[string]any
	if err := json.Unmarshal([]byte(sessionText), &sessionResult); err != nil {
		return fmt.Errorf("解析导入会话响应失败: %w", err)
	}
	sessionID, _ := sessionResult["sessionId"].(string)
	uploadURL, _ := sessionResult["uploadUrl"].(string)
	if sessionID == "" || uploadURL == "" {
		deps.Out.PrintRaw(sessionText)
		return fmt.Errorf("创建导入会话成功但缺少 sessionId 或 uploadUrl")
	}
	deps.Out.PrintInfo(fmt.Sprintf("    会话已创建，sessionId: %s", sessionID))

	deps.Out.PrintInfo("[2/4] 上传文件...")
	if err := httpPutFile(ctx, uploadURL, nil, file.path, file.size); err != nil {
		return fmt.Errorf("文件上传失败 (sessionId=%s): %w", sessionID, err)
	}
	deps.Out.PrintInfo("    文件上传完成")

	deps.Out.PrintInfo("[3/4] 确认导入，启动格式转换...")
	confirmText, err := cfg.callTool(ctx, "confirm_import", map[string]any{"sessionId": sessionID})
	if err != nil {
		return fmt.Errorf("确认导入失败 (sessionId=%s): %w", sessionID, err)
	}
	var confirmResult map[string]any
	if err := json.Unmarshal([]byte(confirmText), &confirmResult); err != nil {
		return fmt.Errorf("解析确认导入响应失败: %w", err)
	}
	taskID, _ := confirmResult["taskId"].(string)
	if taskID == "" {
		deps.Out.PrintRaw(confirmText)
		return fmt.Errorf("确认导入成功但未返回 taskId")
	}
	deps.Out.PrintInfo(fmt.Sprintf("    转换任务已提交，taskId: %s", taskID))

	deps.Out.PrintInfo("[4/4] 等待格式转换完成...")
	result, err := pollImportTask(ctx, taskID, cfg)
	if err != nil {
		var timeoutErr *importPollTimeoutError
		if !errors.As(err, &timeoutErr) {
			return err
		}
		if cfg.timeoutAsResult {
			deps.Out.PrintInfo(timeoutErr.Error())
			_ = deps.Out.PrintJSON(map[string]any{
				"success":      false,
				"timed_out":    true,
				"taskId":       taskID,
				"status":       "processing",
				"next_command": fmt.Sprintf(cfg.nextCommand, taskID),
			})
			return nil
		}
		return fmt.Errorf("%s，请稍后使用 %s 手动查询", timeoutErr.Error(), fmt.Sprintf(cfg.nextCommand, taskID))
	}

	documentURL, _ := result["documentUrl"].(string)
	documentName, _ := result["documentName"].(string)
	documentType, _ := result["documentType"].(string)
	finalResult := map[string]any{
		"success":      true,
		"taskId":       taskID,
		"documentUrl":  documentURL,
		"documentName": documentName,
		"documentType": documentType,
	}
	if cfg.includeNodeID {
		finalResult["nodeId"] = extractNodeIDFromDocURL(documentURL)
	}
	deps.Out.PrintInfo(fmt.Sprintf("导入完成: %s", documentURL))
	_ = deps.Out.PrintJSON(finalResult)
	return nil
}

func runImportGetCommand(cmd *cobra.Command, cfg importFlowConfig) error {
	taskID := mustGetFlag(cmd, "task-id")
	if taskID == "" {
		return fmt.Errorf("flag --task-id is required")
	}
	if deps.Caller.DryRun() {
		deps.Out.PrintKeyValue("操作", cfg.queryOperation)
		deps.Out.PrintKeyValue("任务ID", taskID)
		return nil
	}

	ctx := cmd.Context()
	if ctx == nil {
		ctx = context.Background()
	}
	text, err := cfg.callTool(ctx, "query_import_task", map[string]any{"taskId": taskID})
	if err != nil {
		return err
	}

	var result map[string]any
	if err := json.Unmarshal([]byte(text), &result); err != nil {
		deps.Out.PrintRaw(text)
		return nil
	}
	status, _ := result["status"].(string)
	message, _ := result["message"].(string)
	if strings.EqualFold(status, "completed") {
		if cfg.includeNodeID {
			documentURL, _ := result["documentUrl"].(string)
			result["nodeId"] = extractNodeIDFromDocURL(documentURL)
		}
		_ = deps.Out.PrintJSON(result)
		return nil
	}
	if strings.EqualFold(status, "processing") {
		_ = deps.Out.PrintJSON(result)
		return nil
	}

	_ = deps.Out.PrintJSON(result)
	if message != "" {
		return fmt.Errorf("导入任务失败 (status=%s): %s", status, message)
	}
	return fmt.Errorf("导入任务失败 (status=%s)", status)
}

func pollImportTask(ctx context.Context, taskID string, cfg importFlowConfig) (map[string]any, error) {
	poll := cfg.poll
	if poll.maxPolls <= 0 || poll.interval == nil || poll.wait == nil {
		poll = defaultImportPollPolicy()
	}
	for attempt := 1; attempt <= poll.maxPolls; attempt++ {
		interval := poll.interval(attempt)
		deps.Out.PrintInfo(fmt.Sprintf("    第 %d/%d 次查询，等待 %v ...", attempt, poll.maxPolls, interval))
		if err := poll.wait(ctx, interval); err != nil {
			return nil, fmt.Errorf("导入轮询被取消 (taskId=%s): %w", taskID, err)
		}

		text, err := cfg.callTool(ctx, "query_import_task", map[string]any{"taskId": taskID})
		if err != nil {
			return nil, fmt.Errorf("查询导入任务失败 (taskId=%s): %w", taskID, err)
		}
		var result map[string]any
		if err := json.Unmarshal([]byte(text), &result); err != nil {
			return nil, fmt.Errorf("解析查询结果失败 (taskId=%s): %w", taskID, err)
		}
		status, _ := result["status"].(string)
		switch strings.ToLower(status) {
		case "completed":
			return result, nil
		case "processing":
			continue
		case "failed":
			message, _ := result["message"].(string)
			if message != "" {
				return nil, fmt.Errorf("导入任务失败 (taskId=%s): %s", taskID, message)
			}
			return nil, fmt.Errorf("导入任务失败 (taskId=%s)", taskID)
		}
	}
	return nil, &importPollTimeoutError{taskID: taskID, maxPolls: poll.maxPolls}
}

func extractNodeIDFromDocURL(rawURL string) string {
	parsed, err := url.Parse(rawURL)
	if err != nil || parsed.Path == "" {
		return ""
	}
	nodeID := path.Base(strings.TrimRight(parsed.Path, "/"))
	if nodeID == "." || nodeID == "/" {
		return ""
	}
	return nodeID
}