fiction/backend/app.go

package backend

import (
	"context"
	"encoding/json"
	"fiction/backend/ai"
	"fiction/backend/doubao"
	"fiction/db/model"
	"fiction/db/sqlite"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// App struct
type App struct {
	ctx context.Context
}

// NewApp creates a new App application struct
func NewApp() *App {
	return &App{}
}

func (a *App) Startup(ctx context.Context) {
	a.ctx = ctx
	if err := sqlite.Init("./db/db.db"); err != nil {
		panic(err)
	}
}

func (a *App) Shutdown(ctx context.Context) {
	a.ctx = nil
	sqlite.Stop()
}

// GetNovelList returns the list of txt files in the source directory
func (a *App) GetNovelList() []string {
	var novels []string
	sourceDir := "./source"

	files, err := os.ReadDir(sourceDir)
	if err != nil {
		fmt.Printf("Error reading source directory: %v\n", err)
		return novels
	}

	for _, file := range files {
		if !file.IsDir() && filepath.Ext(file.Name()) == ".txt" {
			novels = append(novels, file.Name())
		}
	}

	return novels
}

// GetNovelContent returns the content of a txt file
func (a *App) GetNovelContent(filename string) string {
	sourceDir := "./source"
	filePath := filepath.Join(sourceDir, filename)

	content, err := os.ReadFile(filePath)
	if err != nil {
		fmt.Printf("Error reading file %s: %v\n", filename, err)
		return "Error reading file"
	}

	return string(content)
}

// GetFangWenContent 获取仿文的内容
func (a *App) GetFangWenContent(filename string) string {
	first := &model.List{}
	sqlite.DB.Model(&model.List{}).
		Where("name = ?", filename).
		First(&first)
	if first.ID > 0 {
		if len(first.Imitation) > 0 {
			return first.Imitation
		}
		if first.Status == 1 {
			return "正在仿写中..."
		}
	}
	return ""
}

// GetAIContent 获取ai对话记录
func (a *App) GetAIContent(filename string) []ai.Message {
	first := &model.List{}
	sqlite.DB.Model(&model.List{}).
		Where("name = ?", filename).
		First(&first)
	if first.ID > 0 {
		if len(first.AI) > 0 {
			resp := make([]ai.Message, 0)
			_ = json.Unmarshal([]byte(first.AI), &resp)
			return resp
		}
	}
	return make([]ai.Message, 0)
}

// GetToken 获取ai token消耗
func (a *App) GetToken(filename string) string {
	first := &model.List{}
	sqlite.DB.Model(&model.List{}).
		Where("name = ?", filename).
		First(&first)
	return first.Token
}

func (a *App) WriteNovelContent(filename string) string {
	content, err := os.ReadFile(filepath.Join("./config", "config.json"))
	if err != nil {
		fmt.Printf("读取提示词文件错误: %v\n", err)
		return "Error reading file"
	}
	var str []string
	if err = json.Unmarshal(content, &str); err != nil {
		fmt.Printf("读取提示词错误: %v\n", err)
		return "Error Unmarshal JSON"
	}

	first := &model.List{}
	sqlite.DB.Model(&model.List{}).
		Where("name = ?", filename).
		First(&first)
	if first.ID == 0 {
		first.Name = filename
		first.Status = 1
		sqlite.DB.Create(&first)
	}

	resp := make([]string, 0)
	msgList, in, out, err := doubao.NewAIClient(false, "", 0).
		RequestAI(
			a.GetModelConfig().Model,
			a.GetModelConfig().Key,
			[]string{
				"本次对话是通过API的方式对话，所以请你不要在回答中掺杂各种图片、修正格式的符号等等，仅以纯文本+换行符的方式回答",
				a.GetNovelContent(filename),
			},
			str,
			func(content string, i int) error {
				processedContent := strings.ReplaceAll(content, "\n\n", "\n")
				resp = append(resp, processedContent)
				return nil
			},
		)
	if err != nil {
		return err.Error()
	}
	marshal, _ := json.Marshal(msgList)

	first.Name = filename
	first.Imitation = resp[len(resp)-1]
	first.AI = string(marshal)
	first.Status = 2
	first.Token = fmt.Sprintf("模型：%v，输入 %d Token，输出 %d Token，预计费用：%.4f元", a.GetModelConfig().Model, in, out, a.GetModelConfig().In/1000000*float64(in)+a.GetModelConfig().Out/1000000*float64(out))
	sqlite.DB.Save(first)

	saveToFile(filename, resp[len(resp)-1])

	return resp[len(resp)-1]
}

func saveToFile(name string, content string) {
	// 获取当前时间戳（精确到秒）
	timestamp := time.Now().Format("01-02 15-04-05") // 格式化为 YYYYMMDDHHMMSS

	// 构建输出文件名：原文件名 + 时间戳 + .txt
	baseName := strings.TrimSuffix(name, filepath.Ext(name))
	outputFileName := fmt.Sprintf("%s_%s.txt", baseName, timestamp)

	outputPath := filepath.Join("./generate", outputFileName)
	if err := os.WriteFile(outputPath, []byte(content), 0644); err != nil {
		fmt.Printf("写入文件失败: %v\n", err)
	}
}