// aes.go
package crypt

import (
	"context"
	"fmt"

	"github.com/sashabaranov/go-openai"
	"github.com/sashabaranov/go-openai/jsonschema"
)

// ========== AES-CBC 加密工具 ==========
type AESCBCEncryptTool struct{}

func (t *AESCBCEncryptTool) Name() string { return "aes_cbc_encrypt" }

func (t *AESCBCEncryptTool) Description() string {
	return "AES-CBC模式加密实现指南。当文档要求使用AES-CBC模式对请求体或敏感字段进行加密时使用，包含PKCS7填充、IV生成和处理"
}

func (t *AESCBCEncryptTool) GetFunctionDefinition() openai.FunctionDefinition {
	return openai.FunctionDefinition{
		Name:        t.Name(),
		Description: t.Description(),
		Parameters: jsonschema.Definition{
			Type: jsonschema.Object,
			Properties: map[string]jsonschema.Definition{
				"key_length": {
					Type:        jsonschema.Integer,
					Description: "密钥长度，可选 16(AES-128), 24(AES-192), 32(AES-256)，默认 32",
					Enum:        []string{"16", "24", "32"},
				},
				"encoding": {
					Type:        jsonschema.String,
					Description: "输出编码方式，base64 或 hex，默认 base64",
					Enum:        []string{"base64", "hex"},
				},
			},
			Required: []string{},
		},
	}
}

func (t *AESCBCEncryptTool) GetDetail(ctx context.Context, params map[string]string) (string, error) {
	keyLength := "32"
	if v, ok := params["key_length"]; ok && v != "" {
		keyLength = v
	}
	encoding := "base64"
	if v, ok := params["encoding"]; ok && v != "" {
		encoding = v
	}

	aesType := "256"
	if keyLength == "16" {
		aesType = "128"
	} else if keyLength == "24" {
		aesType = "192"
	}

	return fmt.Sprintf(`
### AES-CBC 加密完整实现指南

**适用场景**：文档要求使用 AES-CBC 模式进行对称加密

**配置参数**：
- 密钥长度: %s 字节 (AES-%s)
- 编码方式: %s

**完整代码模板**：

`+"```go"+`
package crypto

import (
    "crypto/aes"
    "crypto/cipher"
    "crypto/rand"
    "encoding/base64"
    "encoding/hex"
    "fmt"
    "io"
    "bytes"
)

// AESCBCEncrypt AES-CBC模式加密
func AESCBCEncrypt(plaintext []byte, key []byte) (string, error) {
    block, err := aes.NewCipher(key)
    if err != nil {
        return "", fmt.Errorf("创建AES cipher失败: %%v", err)
    }

    padded := pkcs7Padding(plaintext, aes.BlockSize)

    iv := make([]byte, aes.BlockSize)
    if _, err := io.ReadFull(rand.Reader, iv); err != nil {
        return "", fmt.Errorf("生成IV失败: %%v", err)
    }

    mode := cipher.NewCBCEncrypter(block, iv)
    ciphertext := make([]byte, len(padded))
    mode.CryptBlocks(ciphertext, padded)

    result := append(iv, ciphertext...)
    if "%s" == "hex" {
        return hex.EncodeToString(result), nil
    }
    return base64.StdEncoding.EncodeToString(result), nil
}

// AESCBCDecrypt AES-CBC模式解密
func AESCBCDecrypt(encryptedData string, key []byte) ([]byte, error) {
    var data []byte
    var err error
    
    if "%s" == "hex" {
        data, err = hex.DecodeString(encryptedData)
    } else {
        data, err = base64.StdEncoding.DecodeString(encryptedData)
    }
    if err != nil {
        return nil, fmt.Errorf("解码失败: %%v", err)
    }

    if len(data) < aes.BlockSize {
        return nil, fmt.Errorf("数据长度不足")
    }
    iv := data[:aes.BlockSize]
    ciphertext := data[aes.BlockSize:]

    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, fmt.Errorf("创建AES cipher失败: %%v", err)
    }

    mode := cipher.NewCBCDecrypter(block, iv)
    plaintext := make([]byte, len(ciphertext))
    mode.CryptBlocks(plaintext, ciphertext)

    plaintext, err = pkcs7UnPadding(plaintext)
    if err != nil {
        return nil, fmt.Errorf("去除填充失败: %%v", err)
    }

    return plaintext, nil
}

// pkcs7Padding PKCS7填充
func pkcs7Padding(data []byte, blockSize int) []byte {
    padding := blockSize - len(data)%%blockSize
    padText := bytes.Repeat([]byte{byte(padding)}, padding)
    return append(data, padText...)
}

// pkcs7UnPadding 去除PKCS7填充
func pkcs7UnPadding(data []byte) ([]byte, error) {
    length := len(data)
    if length == 0 {
        return nil, fmt.Errorf("数据为空")
    }
    padding := int(data[length-1])
    if padding > length {
        return nil, fmt.Errorf("填充数据无效")
    }
    return data[:length-padding], nil
}
`+"```"+`

**注意事项**：
- 密钥长度必须是 16/24/32 字节
- IV 必须随机生成且每次不同
- IV 随密文一起传输
`, keyLength, aesType, encoding, encoding, encoding), nil
}

func (t *AESCBCEncryptTool) Execute(ctx context.Context, params map[string]string) (string, error) {
	return t.GetDetail(ctx, params)
}

// ========== AES-ECB 加密工具 ==========
type AESECBEncryptTool struct{}

func (t *AESECBEncryptTool) Name() string { return "aes_ecb_encrypt" }

func (t *AESECBEncryptTool) Description() string {
	return "AES-ECB模式加密实现指南。⚠️注意ECB模式不安全，仅用于兼容老旧系统"
}

func (t *AESECBEncryptTool) GetFunctionDefinition() openai.FunctionDefinition {
	return openai.FunctionDefinition{
		Name:        t.Name(),
		Description: t.Description(),
		Parameters: jsonschema.Definition{
			Type: jsonschema.Object,
			Properties: map[string]jsonschema.Definition{
				"key_length": {
					Type:        jsonschema.Integer,
					Description: "密钥长度，可选 16(AES-128), 24(AES-192), 32(AES-256)，默认 32",
					Enum:        []string{"16", "24", "32"},
				},
				"encoding": {
					Type:        jsonschema.String,
					Description: "输出编码方式，base64 或 hex，默认 base64",
					Enum:        []string{"base64", "hex"},
				},
			},
			Required: []string{},
		},
	}
}

func (t *AESECBEncryptTool) GetDetail(ctx context.Context, params map[string]string) (string, error) {
	keyLength := "32"
	if v, ok := params["key_length"]; ok && v != "" {
		keyLength = v
	}
	encoding := "base64"
	if v, ok := params["encoding"]; ok && v != "" {
		encoding = v
	}

	return fmt.Sprintf(`
### AES-ECB 加密完整实现指南

**配置参数**：
- 密钥长度: %s 字节
- 编码方式: %s

**完整代码模板**：

`+"```go"+`
package crypto

import (
    "crypto/aes"
    "encoding/base64"
    "encoding/hex"
    "fmt"
    "bytes"
)

// AESECBEncrypt AES-ECB模式加密
func AESECBEncrypt(plaintext []byte, key []byte, encoding string) (string, error) {
    block, err := aes.NewCipher(key)
    if err != nil {
        return "", fmt.Errorf("创建AES cipher失败: %%v", err)
    }

    padded := pkcs7Padding(plaintext, aes.BlockSize)

    ciphertext := make([]byte, len(padded))
    for i := 0; i < len(padded); i += aes.BlockSize {
        block.Encrypt(ciphertext[i:i+aes.BlockSize], padded[i:i+aes.BlockSize])
    }

    if encoding == "hex" {
        return hex.EncodeToString(ciphertext), nil
    }
    return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// AESECBDecrypt AES-ECB模式解密
func AESECBDecrypt(encryptedData string, key []byte, encoding string) ([]byte, error) {
    var ciphertext []byte
    var err error
    
    if encoding == "hex" {
        ciphertext, err = hex.DecodeString(encryptedData)
    } else {
        ciphertext, err = base64.StdEncoding.DecodeString(encryptedData)
    }
    if err != nil {
        return nil, fmt.Errorf("解码失败: %%v", err)
    }

    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, fmt.Errorf("创建AES cipher失败: %%v", err)
    }

    if len(ciphertext)%%aes.BlockSize != 0 {
        return nil, fmt.Errorf("密文长度不是块大小的整数倍")
    }

    plaintext := make([]byte, len(ciphertext))
    for i := 0; i < len(ciphertext); i += aes.BlockSize {
        block.Decrypt(plaintext[i:i+aes.BlockSize], ciphertext[i:i+aes.BlockSize])
    }

    plaintext, err = pkcs7UnPadding(plaintext)
    if err != nil {
        return nil, fmt.Errorf("去除填充失败: %%v", err)
    }

    return plaintext, nil
}
`+"```"+`

`, keyLength, encoding, keyLength), nil
}

func (t *AESECBEncryptTool) Execute(ctx context.Context, params map[string]string) (string, error) {
	return t.GetDetail(ctx, params)
}