/**
 * @author: dn-jinmin/dn-jinmin
 * @doc:
 */

package book

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"imooc.com/aicode/practice/book/pkg"
	"imooc.com/aicode/practice/book/prompts"
	"os"
	"strings"
)

type BookGPT struct {
	gpt *pkg.Chat
}

func NewBookGPT() *BookGPT {
	return &BookGPT{gpt: pkg.NewChat()}
}

// content \ count \ lastContent

// 根据标题，生成文章的章节
//   根据文章生成章节数
//   根据主题，章节 -》生成摘要
//	 根据文章主题，标题生成对应的内容  【分批次执行】
//   最后根据文章主题、标题生成总结

// 根据提示词生成 内容

func (b *BookGPT) Gen(title, filepath string) (err error) {
	ctx := context.Background()

	var content strings.Builder
	defer func() {
		if e := b.SaveContent(content.String(), filepath); e != nil || err != nil {
			if err == nil {
				err = e
				return
			}
			err = errors.New(fmt.Sprintf("gen err %v, save file err %v", err, e))
		}
	}()

	content.Write([]byte(fmt.Sprintf("主题: %s", title)))

	// 生成章节数
	sections, err := b.Section(ctx, title)
	if err != nil {
		return err
	}
	fmt.Println("生成章节成功：", sections)

	// 章节的整体字符串
	var sectionStr strings.Builder
	for i, _ := range sections {
		sectionStr.Write([]byte(sections[i].Title))
		if len(sections)-1 > i {
			sectionStr.Write([]byte(","))
		}
	}

	// 先生成摘要
	abstract, err := b.Abstract(ctx, title, sectionStr.String())
	if err != nil {
		return err
	}
	fmt.Println("生成摘要成功：", abstract.Content)
	content.Write([]byte(fmt.Sprintf("\n\n%s", abstract.Content)))

	// 先生成摘要
	sectionContents, err := b.SectionContent(ctx, title, sections)
	if err != nil {
		return err
	}
	fmt.Println("生成章节内容成功：", len(sectionContents))
	for i, _ := range sectionContents {
		content.Write([]byte(fmt.Sprintf("\n\n%s", sectionContents[i].Content)))
	}

	// 总结
	summary, err := b.Summary(ctx, title, sectionStr.String())
	if err != nil {
		return err
	}
	fmt.Println("生成总结成功：", summary.Content)
	content.Write([]byte(fmt.Sprintf("\n\n%s", summary.Content)))

	return nil
}

// 根据主题，章节 -》生成摘要
func (b *BookGPT) Abstract(ctx context.Context, title string, sectionStr string) (*BodyContent, error) {
	prompt := fmt.Sprintf("主题: %s \n 所有章节: %s", title, sectionStr)
	return b.bodyContent(ctx, prompts.Abstract, prompt)
}

// 最后根据文章主题、标题生成总结
func (b *BookGPT) Summary(ctx context.Context, title string, sectionStr string) (*BodyContent, error) {
	prompt := fmt.Sprintf("主题: %s \n 所有章节: %s", title, sectionStr)
	return b.bodyContent(ctx, prompts.Summary, prompt)
}

func (b *BookGPT) SectionContent(ctx context.Context, title string, sections []*Section) ([]*BodyContent, error) {
	res := make([]*BodyContent, 0, len(sections))

	for i, _ := range sections {
		fmt.Println("生成章节", sections[i].Title)

		// 直接通过gpt生成内容
		prompt := fmt.Sprintf("主题: %s\n 章节：%s \n内容字数要求: %d", title, sections[i].Title, sections[i].Count)
		bodyContent, err := b.bodyContent(ctx, prompts.SectionContent, prompt)
		if err != nil {
			return nil, err
		}

		res = append(res, bodyContent)

		// 判断生成的内容是不是完成了
		count := bodyContent.Count
		for sections[i].Count > count {
			// 接着完善文章内容
			fmt.Println("继续完善章节: ", sections[i].Title, " count : ", sections[i].Count, " new count : ", count)

			perfectPrompt := fmt.Sprintf("主题：%s\n章节：%s\n之前的内容：%s\n内容字数要求：%v", title, sections[i].Title,
				bodyContent.LastContent, sections[i].Count-count)
			bodyContent, err = b.bodyContent(ctx, prompts.PerfectContent, perfectPrompt)
			if err != nil {
				return nil, err
			}
			res = append(res, bodyContent)

			count += bodyContent.Count
		}

		fmt.Println("生成章节", sections[i].Title, " 完成")
	}

	return res, nil
}

// 生成文本内容的方法
func (b *BookGPT) bodyContent(ctx context.Context, system, prompt string) (*BodyContent, error) {
	resp, err := b.gpt.GenResponse(ctx, system, prompt)
	if err != nil {
		return nil, err
	}

	var res BodyContent
	if err := json.Unmarshal([]byte(resp), &res); err != nil {

		fmt.Println("bodyContent ", err, " resp ", resp)
		return nil, err
	}

	return &res, nil
}

// 根据文章标题 生成章节数
func (b *BookGPT) Section(ctx context.Context, title string) ([]*Section, error) {
	resp, err := b.gpt.GenResponse(ctx, prompts.Section, title)
	if err != nil {
		return nil, err
	}

	var section []*Section
	if err := json.Unmarshal([]byte(resp), &section); err != nil {
		fmt.Println("Section ", err, " resp ", resp)
		return nil, err
	}

	return section, nil
}

func (b *BookGPT) SaveContent(content, filePath string) error {
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		return os.WriteFile(filePath, []byte(content), 0644)
	}

	file, err := os.OpenFile(filePath, os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.WriteString(content)
	return err
}