package rsa

import (
	"fmt"
	"io"
	"os"

	"archgrid.xyz/ag/tools/types"
	"archgrid.xyz/ag/toolsbox/encryption/rsa"
	"archgrid.xyz/ag/toolsbox/serialize/base64"
	"archgrid.xyz/ag/toolsbox/serialize/hex"
	"github.com/spf13/cobra"
)

var (
	rsaEncryptFromFile   bool
	rsaEncryptToFile     bool
	rsaResultOutput      types.ResultEncoding = types.ResultInHex
	rsaEncryptInputFile  string
	rsaEncryptOutputFile string
)

var rsaEncryptCmd = &cobra.Command{
	Use:   "encrypt",
	Short: `使用RSA算法对给定的内容进行加密`,
	Long:  `使用RSA算法对给定的内容进行加密，生成RSA密钥。如果指定了输入文件，则从输入文件中读取内容进行加密，否则从命令行读取内容进行加密。如果指定了输出文件，则将加密后的内容写入到输出文件中，否则将加密后的内容输出到命令行。`,
	Args:  cobra.MaximumNArgs(1),
	Run:   rsaEncryptExecute,
}

func rsaEncryptExecute(cmd *cobra.Command, args []string) {
	var (
		content []byte
		cert    []byte
		keys    *rsa.KeyPair
		err     error
	)
	if len(publicKeyFile) > 0 {
		file, err := os.Open(publicKeyFile)
		if err != nil {
			fmt.Printf("无法打开RSA公钥文件：%s\n", err)
			return
		}
		cert, err = io.ReadAll(file)
		if err != nil {
			fmt.Printf("无法读取RSA公钥文件：%s\n", err)
			return
		}
	} else if len(privateKeyFile) > 0 {
		file, err := os.Open(privateKeyFile)
		if err != nil {
			fmt.Printf("无法打开RSA私钥文件：%s\n", err)
			return
		}
		cert, err = io.ReadAll(file)
		if err != nil {
			fmt.Printf("无法读取RSA私钥文件：%s\n", err)
			return
		}
	} else {
		fmt.Println("必须指定RSA公钥文件或RSA私钥文件。")
		return
	}
	keys, err = rsa.NewFromPEM(cert)
	if err != nil {
		fmt.Printf("无法解析RSA密钥：%s\n", err)
		return
	}
	if rsaEncryptFromFile {
		if len(rsaEncryptInputFile) == 0 {
			fmt.Println("必须指定输入文件。")
			return
		}
		file, err := os.Open(rsaEncryptInputFile)
		if err != nil {
			fmt.Printf("无法打开输入文件：%s\n", err)
			return
		}
		content, err = io.ReadAll(file)
		if err != nil {
			fmt.Printf("无法读取输入文件：%s\n", err)
			return
		}
	} else {
		if len(args) == 0 {
			fmt.Println("必须指定要加密的内容。")
			return
		}
		content = []byte(args[0])
	}
	encrypted, err := keys.Encrypt(content)
	if err != nil {
		fmt.Printf("无法加密内容：%s\n", err)
		return
	}
	if len(rsaEncryptOutputFile) > 0 {
		file, err := os.OpenFile(rsaEncryptOutputFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
		if err != nil {
			fmt.Printf("无法创建RSA加密文件：%s\n", err)
			return
		}
		_, err = file.Write(encrypted)
		if err != nil {
			fmt.Printf("无法写入RSA加密文件：%s\n", err)
			return
		}
		fmt.Printf("已将加密结果写入到文件：%s\n", rsaEncryptOutputFile)
	} else {
		switch rsaResultOutput {
		case types.ResultInBase64:
			fmt.Printf("加密结果：%s\n", base64.ToBase64(encrypted))
		case types.ResultInHex:
			fallthrough
		default:
			fmt.Printf("加密结果：%s\n", hex.ToHex(encrypted))
		}
	}
}

func init() {
	rsaEncryptCmd.PersistentFlags().BoolVarP(&rsaEncryptFromFile, "from-file", "f", false, "从输入文件中读取内容进行加密。")
	rsaEncryptCmd.PersistentFlags().BoolVarP(&rsaEncryptToFile, "to-file", "t", false, "将加密后的内容写入到输出文件中。")
	rsaEncryptCmd.PersistentFlags().StringVarP(&rsaEncryptInputFile, "input-file", "i", "", "输入文件。")
	rsaEncryptCmd.PersistentFlags().StringVarP(&rsaEncryptOutputFile, "output-file", "o", "", "输出文件。")
	rsaEncryptCmd.PersistentFlags().VarP(&rsaResultOutput, "output-encode", "e", "加密结果输出格式，可选值有：hex，base64，缺省值为hex。对输出到文件的方式无效。")
	rsaCmd.AddCommand(rsaEncryptCmd)
}