package utils

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"image"
	"image/png"
	"io"
)

// CharacterCardV2 SillyTavern 角色卡 V2 格式
type CharacterCardV2 struct {
	Spec        string              `json:"spec"`
	SpecVersion string              `json:"spec_version"`
	Data        CharacterCardV2Data `json:"data"`
}

type CharacterCardV2Data struct {
	Name                    string                 `json:"name"`
	Description             string                 `json:"description"`
	Personality             string                 `json:"personality"`
	Scenario                string                 `json:"scenario"`
	FirstMes                string                 `json:"first_mes"`
	MesExample              string                 `json:"mes_example"`
	CreatorNotes            string                 `json:"creator_notes"`
	SystemPrompt            string                 `json:"system_prompt"`
	PostHistoryInstructions string                 `json:"post_history_instructions"`
	Tags                    []string               `json:"tags"`
	Creator                 string                 `json:"creator"`
	CharacterVersion        string                 `json:"character_version"`
	AlternateGreetings      []string               `json:"alternate_greetings"`
	Extensions              map[string]interface{} `json:"extensions"`
}

// ExtractCharacterFromPNG 从 PNG 图片中提取角色卡数据
func ExtractCharacterFromPNG(pngData []byte) (*CharacterCardV2, error) {
	reader := bytes.NewReader(pngData)

	// 验证 PNG 格式（解码但不保存图片）
	_, err := png.Decode(reader)
	if err != nil {
		return nil, errors.New("无效的 PNG 文件")
	}

	// 重新读取以获取 tEXt chunks
	reader.Seek(0, 0)

	// 查找 tEXt chunk 中的 "chara" 字段
	charaJSON, err := extractTextChunk(reader, "chara")
	if err != nil {
		return nil, errors.New("PNG 中没有找到角色卡数据")
	}

	// 尝试 Base64 解码
	decodedJSON, err := base64.StdEncoding.DecodeString(charaJSON)
	if err != nil {
		// 如果不是 Base64，直接使用原始 JSON
		decodedJSON = []byte(charaJSON)
	}

	// 解析 JSON
	var card CharacterCardV2
	err = json.Unmarshal(decodedJSON, &card)
	if err != nil {
		return nil, errors.New("解析角色卡数据失败: " + err.Error())
	}

	return &card, nil
}

// extractTextChunk 从 PNG 中提取指定 key 的 tEXt chunk
func extractTextChunk(r io.Reader, key string) (string, error) {
	// 跳过 PNG signature (8 bytes)
	signature := make([]byte, 8)
	if _, err := io.ReadFull(r, signature); err != nil {
		return "", err
	}

	// 验证 PNG signature
	expectedSig := []byte{137, 80, 78, 71, 13, 10, 26, 10}
	if !bytes.Equal(signature, expectedSig) {
		return "", errors.New("invalid PNG signature")
	}

	// 读取所有 chunks
	for {
		// 读取 chunk length (4 bytes)
		lengthBytes := make([]byte, 4)
		if _, err := io.ReadFull(r, lengthBytes); err != nil {
			if err == io.EOF {
				break
			}
			return "", err
		}
		length := uint32(lengthBytes[0])<<24 | uint32(lengthBytes[1])<<16 |
			uint32(lengthBytes[2])<<8 | uint32(lengthBytes[3])

		// 读取 chunk type (4 bytes)
		chunkType := make([]byte, 4)
		if _, err := io.ReadFull(r, chunkType); err != nil {
			return "", err
		}

		// 读取 chunk data
		data := make([]byte, length)
		if _, err := io.ReadFull(r, data); err != nil {
			return "", err
		}

		// 读取 CRC (4 bytes)
		crc := make([]byte, 4)
		if _, err := io.ReadFull(r, crc); err != nil {
			return "", err
		}

		// 检查是否是 tEXt chunk
		if string(chunkType) == "tEXt" {
			// tEXt chunk 格式: keyword\0text
			nullIndex := bytes.IndexByte(data, 0)
			if nullIndex == -1 {
				continue
			}

			keyword := string(data[:nullIndex])
			text := string(data[nullIndex+1:])

			if keyword == key {
				return text, nil
			}
		}

		// IEND chunk 表示结束
		if string(chunkType) == "IEND" {
			break
		}
	}

	return "", errors.New("text chunk not found")
}

// EmbedCharacterToPNG 将角色卡数据嵌入到 PNG 图片中
func EmbedCharacterToPNG(img image.Image, card *CharacterCardV2) ([]byte, error) {
	// 序列化角色卡数据
	cardJSON, err := json.Marshal(card)
	if err != nil {
		return nil, err
	}

	// Base64 编码
	encodedJSON := base64.StdEncoding.EncodeToString(cardJSON)

	// 创建一个 buffer 来写入 PNG
	var buf bytes.Buffer

	// 写入 PNG signature
	buf.Write([]byte{137, 80, 78, 71, 13, 10, 26, 10})

	// 编码原始图片到临时 buffer
	var imgBuf bytes.Buffer
	if err := png.Encode(&imgBuf, img); err != nil {
		return nil, err
	}

	// 跳过原始 PNG 的 signature
	imgData := imgBuf.Bytes()[8:]

	// 将原始图片的 chunks 复制到输出，在 IEND 之前插入 tEXt chunk
	r := bytes.NewReader(imgData)

	for {
		// 读取 chunk length
		lengthBytes := make([]byte, 4)
		if _, err := io.ReadFull(r, lengthBytes); err != nil {
			if err == io.EOF {
				break
			}
			return nil, err
		}
		length := uint32(lengthBytes[0])<<24 | uint32(lengthBytes[1])<<16 |
			uint32(lengthBytes[2])<<8 | uint32(lengthBytes[3])

		// 读取 chunk type
		chunkType := make([]byte, 4)
		if _, err := io.ReadFull(r, chunkType); err != nil {
			return nil, err
		}

		// 读取 chunk data
		data := make([]byte, length)
		if _, err := io.ReadFull(r, data); err != nil {
			return nil, err
		}

		// 读取 CRC
		crc := make([]byte, 4)
		if _, err := io.ReadFull(r, crc); err != nil {
			return nil, err
		}

		// 如果是 IEND chunk，先写入 tEXt chunk
		if string(chunkType) == "IEND" {
			// 写入 tEXt chunk
			writeTextChunk(&buf, "chara", encodedJSON)
		}

		// 写入原始 chunk
		buf.Write(lengthBytes)
		buf.Write(chunkType)
		buf.Write(data)
		buf.Write(crc)

		if string(chunkType) == "IEND" {
			break
		}
	}

	return buf.Bytes(), nil
}

// writeTextChunk 写入 tEXt chunk
func writeTextChunk(w io.Writer, keyword, text string) error {
	data := append([]byte(keyword), 0)
	data = append(data, []byte(text)...)

	// 写入 length
	length := uint32(len(data))
	lengthBytes := []byte{
		byte(length >> 24),
		byte(length >> 16),
		byte(length >> 8),
		byte(length),
	}
	w.Write(lengthBytes)

	// 写入 type
	w.Write([]byte("tEXt"))

	// 写入 data
	w.Write(data)

	// 计算并写入 CRC
	crcData := append([]byte("tEXt"), data...)
	crc := calculateCRC(crcData)
	crcBytes := []byte{
		byte(crc >> 24),
		byte(crc >> 16),
		byte(crc >> 8),
		byte(crc),
	}
	w.Write(crcBytes)

	return nil
}

// calculateCRC 计算 CRC32
func calculateCRC(data []byte) uint32 {
	crc := uint32(0xFFFFFFFF)

	for _, b := range data {
		crc ^= uint32(b)
		for i := 0; i < 8; i++ {
			if crc&1 != 0 {
				crc = (crc >> 1) ^ 0xEDB88320
			} else {
				crc >>= 1
			}
		}
	}

	return crc ^ 0xFFFFFFFF
}

// ParseCharacterCardJSON 解析 JSON 格式的角色卡
func ParseCharacterCardJSON(jsonData []byte) (*CharacterCardV2, error) {
	var card CharacterCardV2
	err := json.Unmarshal(jsonData, &card)
	if err != nil {
		return nil, errors.New("解析角色卡 JSON 失败: " + err.Error())
	}

	return &card, nil
}