package pvf

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"strings"
	"unicode/utf16"
)

// 本文件补齐 companionpvf 需要的两个字符串池追加能力。
//
// 背景：支援兵的 AIC 覆盖要改写 `aicharacter/**.aic` 里的技能/按键引用，而
// 这些引用是字符串池下标。如果目标字符串在原始 Script.pvf 的池子里不存在
// （例如新的 key 路径），就必须**追加**到池子尾部，并返回新的魔法偏移。
// 追加只影响池子尾部，已有条目的偏移不变，因此文件表与 body 全部保持原样。
//
// 两个方法的语义与 Server-90 完全一致：
//   - EnsureANSIString ：往 strA 追加 NUL 结尾的 ANSI 串，返回 start<<1（偶数=ANSI）
//   - EnsureUnicodeString：往 strW 追加 UTF-16LE 串，返回 start|1（奇数=UTF-16）
// 返回值是客户端/读取器共用的"魔法偏移"编码，奇偶性就是宽度选择器。

// ensureANSIStringIndex 缓存 ANSI 池里"已有字符串 -> 偏移"的映射，
// 避免每次追加都线性扫描整个池子。
func (a *Archive) ensureANSIStringOffsets() map[string]int {
	if a.ansiOffsets != nil {
		return a.ansiOffsets
	}
	a.ansiOffsets = make(map[string]int)
	for start := 0; start < len(a.strA); {
		endRel := bytes.IndexByte(a.strA[start:], 0)
		if endRel < 0 {
			break
		}
		end := start + endRel
		stringValue := string(a.strA[start:end])
		if _, exists := a.ansiOffsets[stringValue]; !exists {
			a.ansiOffsets[stringValue] = start << 1
		}
		start = end + 1
	}
	return a.ansiOffsets
}

// EnsureANSIString 返回 value 在 ANSI 字符串池中的魔法偏移，必要时追加。
func (a *Archive) EnsureANSIString(value string) (int, error) {
	if a == nil {
		return 0, fmt.Errorf("%w: archive is nil", ErrInvalidArchive)
	}
	if a.format != FormatProtectedNKPI {
		return 0, fmt.Errorf("%w: ANSI append requires %q, got %q", ErrInvalidArchive, FormatProtectedNKPI, a.format)
	}
	if strings.ContainsRune(value, 0) {
		return 0, fmt.Errorf("%w: ANSI string must not contain NUL", ErrInvalidArchive)
	}
	offsets := a.ensureANSIStringOffsets()
	if offset, ok := offsets[value]; ok {
		return offset, nil
	}
	raw := []byte(value)
	if len(a.strA) > math.MaxInt32-len(raw)-1 {
		return 0, fmt.Errorf("%w: ANSI string pool too large", ErrInvalidArchive)
	}
	start := len(a.strA)
	a.strA = append(a.strA, raw...)
	a.strA = append(a.strA, 0)
	offset := start << 1
	offsets[value] = offset
	return offset, nil
}

// RepackRawAndStringPools 同时重建"文件体"和"字符串池"两段，供离线补丁使用。
//
// 与 RepackRaw 的唯一区别在收尾：RepackRaw 直接复制原来的 name 段，而本方法
// 用当前的 a.strA/a.strW 重新编码 name 段。这一点是必需的 —— EnsureANSIString /
// EnsureUnicodeString 只是把新字符串追加到内存池尾部并返回新偏移，如果不把
// 池子写回归档，补丁后的引用就会指向不存在的数据（读出来是空串）。
//
// 只允许 protected NKPI：普通 NKPI 的 name 段布局不同，写回会破坏归档。
func (a *Archive) RepackRawAndStringPools(replacements []RawReplacement) ([]byte, error) {
	if a == nil {
		return nil, fmt.Errorf("%w: archive is nil", ErrInvalidArchive)
	}
	if len(replacements) == 0 {
		return nil, fmt.Errorf("%w: replacement is empty", ErrInvalidArchive)
	}
	if a.format != FormatProtectedNKPI {
		return nil, fmt.Errorf("%w: raw+strings repack requires %q, got %q", ErrInvalidArchive, FormatProtectedNKPI, a.format)
	}
	layout, err := a.layout()
	if err != nil {
		return nil, err
	}

	replacedByIndex := make(map[int][]byte, len(replacements))
	touchedChunks := make(map[int]bool, len(replacements))
	for _, replacement := range replacements {
		idx := a.FindFileIndex(replacement.Path)
		if idx < 0 {
			return nil, fmt.Errorf("%w: %s", ErrFileNotFound, replacement.Path)
		}
		item := a.items[idx]
		if item.chunkIndex < 0 || item.chunkIndex >= len(a.groups) {
			return nil, fmt.Errorf("%w: file %d chunk %d is invalid", ErrInvalidArchive, idx, item.chunkIndex)
		}
		replacedByIndex[idx] = append([]byte(nil), replacement.Data...)
		touchedChunks[item.chunkIndex] = true
	}

	newItems := append([]fileItem(nil), a.items...)
	chunkFiles, err := a.filesByChunk()
	if err != nil {
		return nil, err
	}
	body, newGroups, err := a.repackBody(layout, chunkFiles, touchedChunks, replacedByIndex, newItems)
	if err != nil {
		return nil, err
	}
	if len(body) > math.MaxInt32 {
		return nil, fmt.Errorf("%w: body too large: %d", ErrInvalidArchive, len(body))
	}

	nameSection, err := a.encodeProtectedNameSection(layout, a.strA, a.strW)
	if err != nil {
		return nil, err
	}
	header := a.header
	header.bodySize = len(body)
	header.groupCount = len(newGroups)
	header.nameSize = len(nameSection)

	out := a.encodeHeader(header)
	out = append(out, encodeFileTable(newItems)...)
	out = append(out, a.data[layout.hashOffset:layout.hashOffset+layout.hashSize]...)
	out = append(out, nameSection...)
	out = append(out, a.encodeGroups(newGroups)...)
	out = append(out, body...)
	return out, nil
}

// EnsureUnicodeString 返回 value 在 UTF-16 字符串池中的魔法偏移，必要时追加。
func (a *Archive) EnsureUnicodeString(value string) (int, error) {
	if a == nil || a.format != FormatProtectedNKPI || strings.ContainsRune(value, 0) {
		return 0, fmt.Errorf("%w: Unicode append requires protected NKPI and a NUL-free string", ErrInvalidArchive)
	}
	units := utf16.Encode([]rune(value))
	raw := make([]byte, (len(units)+1)*2)
	for i, unit := range units {
		binary.LittleEndian.PutUint16(raw[i*2:], unit)
	}
	for start := 0; start < len(a.strW); {
		end := start
		for end+1 < len(a.strW) && (a.strW[end] != 0 || a.strW[end+1] != 0) {
			end += 2
		}
		if end+2 > len(a.strW) {
			break
		}
		if bytes.Equal(a.strW[start:end+2], raw) {
			return start | 1, nil
		}
		start = end + 2
	}
	if len(a.strW)%2 != 0 || len(a.strW) > math.MaxInt32-len(raw) {
		return 0, fmt.Errorf("%w: invalid or oversized Unicode string pool", ErrInvalidArchive)
	}
	start := len(a.strW)
	a.strW = append(a.strW, raw...)
	return start | 1, nil
}