package fqapi

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"encoding/binary"
	"fmt"
)

// ============ SpadeA（黑桃A）PlayAuth 解码 ============
//
// 番茄短剧/汽水音乐同源的本地 DRM：video_model 每个清晰度的 spade_a
// （PlayAuth）经混淆变换直接还原 16 字节 AES-128 内容密钥（32 位 hex 字符串）。
// 算法逆向自 SodaMusic v3.5.1（qishui-decrypt-tool），与番茄短剧样本实测一致。

func bitCount32(v uint32) int {
	c := v
	c -= (c >> 1) & 0x55555555
	c = (c & 0x33333333) + ((c >> 2) & 0x33333333)
	return int(((c + (c>>4)) & 0x0f0f0f0f) * 0x01010101 >> 24)
}

func decodeBase36Byte(v byte) int {
	switch {
	case v >= '0' && v <= '9':
		return int(v - '0')
	case v >= 'a' && v <= 'z':
		return int(v-'a') + 10
	}
	return 0xff
}

// decodeSpadeInner 逐字节去混淆：与前一字节异或后减去索引 popcount 与常数
func decodeSpadeInner(data []byte) []byte {
	result := make([]byte, len(data))
	buff := append([]byte{0xfa, 0x55}, data...)
	for i := 0; i < len(data); i++ {
		raw := int(data[i]^buff[i]) - bitCount32(uint32(i)) - 21
		if raw >= 0 {
			result[i] = byte(raw)
		} else {
			result[i] = byte(((raw % 255) + 255) % 255)
		}
	}
	return result
}

// DecodeSpadeKey 从 spade_a（PlayAuth base64）解出 16 字节 AES 内容密钥
func DecodeSpadeKey(playAuthB64 string) ([]byte, error) {
	raw, err := base64.StdEncoding.DecodeString(playAuthB64)
	if err != nil {
		return nil, fmt.Errorf("spade_a base64 解码失败: %w", err)
	}
	if len(raw) < 4 {
		return nil, fmt.Errorf("spade_a 长度不足: %d", len(raw))
	}
	padding := int(raw[0]^raw[1]^raw[2]) - 48
	if padding < 0 || padding >= len(raw)-3 {
		return nil, fmt.Errorf("spade_a padding 非法: %d", padding)
	}
	body := raw[1 : len(raw)-padding]
	tmp := decodeSpadeInner(body)
	keyLen := decodeBase36Byte(tmp[0])
	if keyLen == 0xff {
		return nil, fmt.Errorf("spade_a key 长度字符非法: %q", tmp[0])
	}
	end := len(body) - keyLen
	if end < 1 || end > len(tmp) {
		return nil, fmt.Errorf("spade_a key 范围非法: end=%d len=%d", end, len(tmp))
	}
	keyHex := string(tmp[1:end])
	if len(keyHex) != 32 {
		return nil, fmt.Errorf("spade_a key 长度异常: %d", len(keyHex))
	}
	key := make([]byte, 16)
	for i := 0; i < 16; i++ {
		if _, err := fmt.Sscanf(keyHex[i*2:i*2+2], "%02x", &key[i]); err != nil {
			return nil, fmt.Errorf("spade_a key hex 非法: %s", keyHex)
		}
	}
	return key, nil
}

// ============ CENC (cenc-aes-ctr) MP4 解密 ============
//
// 容器为标准 CENC MP4：moov 明文；加密 sample entry（encv/enca）+ senc box
// （version 0，每个 sample 8 字节 IV，整样本 AES-128-CTR，无 subsample）。
// 解密后把 entry 类型改回原始 codec（sinf/frma 指示），得到可播放文件。

type mp4Box struct {
	typ         string
	contentStar int
	end         int
}

// walkBoxes 迭代 [start,end) 内的顶层 box，fn 返回 false 时停止迭代
func walkBoxes(data []byte, start, end int, fn func(mp4Box) bool) error {
	p := start
	for p+8 <= end {
		size := int(binary.BigEndian.Uint32(data[p : p+4]))
		typ := string(data[p+4 : p+8])
		if size == 1 { // largesize
			if p+16 > end {
				return fmt.Errorf("box %s largesize 越界", typ)
			}
			size = int(binary.BigEndian.Uint64(data[p+8 : p+16]))
			// ponytail: >2GB 文件不支持，截断风险可接受
			if size > len(data) {
				return fmt.Errorf("box %s 过大", typ)
			}
		}
		if size < 8 || p+size > end {
			return fmt.Errorf("box %s 尺寸非法 size=%d at %d", typ, size, p)
		}
		if !fn(mp4Box{typ: typ, contentStar: p + 8, end: p + size}) {
			return nil
		}
		p += size
	}
	return nil
}

// DecryptCENCMP4 原地解密 CENC 加密的 MP4 字节并修复 stsd entry 类型。
// 仅支持 senc version 0 + 无 subsample 标志（番茄实测形态），其他形态返回错误。
func DecryptCENCMP4(data []byte, key []byte) error {
	block, err := aes.NewCipher(key)
	if err != nil {
		return fmt.Errorf("内容密钥非法: %w", err)
	}

	decryptStream := func(ct []byte, iv8 []byte) error {
		if len(iv8) != 8 {
			return fmt.Errorf("senc IV 长度异常: %d", len(iv8))
		}
		iv := make([]byte, 16)
		copy(iv, iv8) // 8 字节 IV 置于高位，低 8 字节为块计数器（CENC 标准）
		cipher.NewCTR(block, iv).XORKeyStream(ct, ct)
		return nil
	}

	// 找 moov 并遍历其下所有 trak，逐 trak 解密
	var trakErr error
	foundMoov := false
	if err := walkBoxes(data, 0, len(data), func(b mp4Box) bool {
		if b.typ != "moov" {
			return true
		}
		foundMoov = true
		_ = walkBoxes(data, b.contentStar, b.end, func(trak mp4Box) bool {
			if trak.typ != "trak" {
				return true
			}
			if err := decryptTrak(data, trak, decryptStream); err != nil {
				trakErr = err
				return false
			}
			return true
		})
		return false
	}); err != nil {
		return err
	}
	if !foundMoov {
		return fmt.Errorf("moov box 未找到")
	}
	return trakErr
}

// decryptTrak 处理单个 trak：解析 sample 表 + senc，逐样本解密，修补 stsd
func decryptTrak(data []byte, trak mp4Box, decryptStream func([]byte, []byte) error) error {
	// 定位 mdia→minf→stbl
	stbl, err := findBox(data, trak, "mdia", "minf", "stbl")
	if err != nil || stbl == nil {
		return err // 无 stbl 的 trak 跳过
	}

	var (
		sizes         []int
		chunkOffsets  []int
		stscEntries   [][2]int // {first_chunk, samples_per_chunk}
		ivs           [][]byte
		stsd          *mp4Box
		haveSizes     bool
		haveOffsets   bool
		haveStsc      bool
		parseErr      error
	)

	if err := walkBoxes(data, stbl.contentStar, stbl.end, func(b mp4Box) bool {
		switch b.typ {
		case "stsd":
			stsd = &b
		case "stsz":
			constSize := int(binary.BigEndian.Uint32(data[b.contentStar+4 : b.contentStar+8]))
			count := int(binary.BigEndian.Uint32(data[b.contentStar+8 : b.contentStar+12]))
			if constSize != 0 {
				for i := 0; i < count; i++ {
					sizes = append(sizes, constSize)
				}
			} else {
				sizes = make([]int, count)
				for i := 0; i < count; i++ {
					sizes[i] = int(binary.BigEndian.Uint32(data[b.contentStar+12+i*4 : b.contentStar+16+i*4]))
				}
			}
			haveSizes = true
		case "stco":
			count := int(binary.BigEndian.Uint32(data[b.contentStar+4 : b.contentStar+8]))
			for i := 0; i < count; i++ {
				chunkOffsets = append(chunkOffsets, int(binary.BigEndian.Uint32(data[b.contentStar+8+i*4:b.contentStar+12+i*4])))
			}
			haveOffsets = true
		case "co64":
			count := int(binary.BigEndian.Uint32(data[b.contentStar+4 : b.contentStar+8]))
			for i := 0; i < count; i++ {
				v := binary.BigEndian.Uint64(data[b.contentStar+8+i*8 : b.contentStar+16+i*8])
				if v > uint64(len(data)) {
					parseErr = fmt.Errorf("co64 偏移越界")
					return false
				}
				chunkOffsets = append(chunkOffsets, int(v))
			}
			haveOffsets = true
		case "stsc":
			count := int(binary.BigEndian.Uint32(data[b.contentStar+4 : b.contentStar+8]))
			for i := 0; i < count; i++ {
				first := int(binary.BigEndian.Uint32(data[b.contentStar+8+i*12 : b.contentStar+12+i*12]))
				spc := int(binary.BigEndian.Uint32(data[b.contentStar+12+i*12 : b.contentStar+16+i*12]))
				stscEntries = append(stscEntries, [2]int{first, spc})
			}
			haveStsc = true
		case "senc":
			verFlags := binary.BigEndian.Uint32(data[b.contentStar : b.contentStar+4])
			if verFlags>>24 != 0 {
				parseErr = fmt.Errorf("senc version %d 不支持", verFlags>>24)
				return false
			}
			if verFlags&0b10 != 0 {
				parseErr = fmt.Errorf("senc subsample 加密模式不支持")
				return false
			}
			count := int(binary.BigEndian.Uint32(data[b.contentStar+4 : b.contentStar+8]))
			p := b.contentStar + 8
			if p+8*count > b.end {
				parseErr = fmt.Errorf("senc IV 表越界")
				return false
			}
			for i := 0; i < count; i++ {
				ivs = append(ivs, data[p+8*i:p+8*i+8])
			}
		}
		return true
	}); err != nil {
		return fmt.Errorf("stbl 解析失败: %w", err)
	}
	if parseErr != nil {
		return parseErr
	}

	if len(ivs) == 0 {
		return nil // 本 trak 未加密
	}
	if !haveSizes || !haveOffsets || !haveStsc {
		return fmt.Errorf("sample 表不完整（stsz=%v stco=%v stsc=%v）", haveSizes, haveOffsets, haveStsc)
	}

	// stsc+stco+stsz 重建每样本绝对偏移（标准 sample table 遍历）
	sampleIdx := 0
	for ci, off := range chunkOffsets {
		spc := stscEntries[len(stscEntries)-1][1]
		for _, e := range stscEntries {
			if e[0] <= ci+1 {
				spc = e[1]
			} else {
				break
			}
		}
		pos := off
		for s := 0; s < spc && sampleIdx < len(sizes); s++ {
			size := sizes[sampleIdx]
			if sampleIdx >= len(ivs) {
				break
			}
			if pos+size > len(data) {
				return fmt.Errorf("sample %d 越界 pos=%d size=%d", sampleIdx, pos, size)
			}
			if err := decryptStream(data[pos:pos+size], ivs[sampleIdx]); err != nil {
				return err
			}
			pos += size
			sampleIdx++
		}
	}

	if stsd != nil {
		patchStsdEntry(data, *stsd)
	}
	return nil
}

// findBox 沿 path 逐层查找唯一子 box；任意一层缺失返回 (nil, nil)
func findBox(data []byte, parent mp4Box, path ...string) (*mp4Box, error) {
	cur := parent
	for _, name := range path {
		var next *mp4Box
		if err := walkBoxes(data, cur.contentStar, cur.end, func(b mp4Box) bool {
			if b.typ == name {
				next = &b
				return false
			}
			return true
		}); err != nil {
			return nil, err
		}
		if next == nil {
			return nil, nil
		}
		cur = *next
	}
	return &cur, nil
}

// patchStsdEntry 把加密 sample entry（encv/enca）的类型改回 sinf/frma 指示的原始 codec。
// 番茄实测单 entry；entry 头部长度因音/视频而异，直接在 entry 字节里定位 "sinf"。
func patchStsdEntry(data []byte, stsd mp4Box) {
	entryCount := int(binary.BigEndian.Uint32(data[stsd.contentStar+4 : stsd.contentStar+8]))
	if entryCount != 1 {
		return // ponytail: 多 entry 形态按需扩展
	}
	entryStart := stsd.contentStar + 8
	entryEnd := stsd.end

	// 在 entry 内定位 "sinf"（box 类型魔数），再验证其 box 头合法性
	sinfOff := bytes.Index(data[entryStart:entryEnd], []byte("sinf"))
	for sinfOff >= 0 {
		boxStart := entryStart + sinfOff - 4 // size 字段起点
		if boxStart >= entryStart+8 {
			size := int(binary.BigEndian.Uint32(data[boxStart : boxStart+4]))
			if size >= 12 && boxStart+size <= entryEnd {
				sinfContent, sinfEnd := boxStart+8, boxStart+size
				frmaOff := bytes.Index(data[sinfContent:sinfEnd], []byte("frma"))
				if frmaOff >= 0 {
					frmaContent := sinfContent + frmaOff + 4
					if frmaContent+4 <= sinfEnd {
						// entry 类型字段在 box 头的 type 位置（entryStart+4）
						copy(data[entryStart+4:entryStart+8], data[frmaContent:frmaContent+4])
						return
					}
				}
			}
		}
		next := bytes.Index(data[entryStart+sinfOff+1:entryEnd], []byte("sinf"))
		if next < 0 {
			return
		}
		sinfOff = sinfOff + 1 + next
	}
}