// Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package nodes

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"github.com/TeaOSLab/EdgeNode/internal/compressions"
	"github.com/TeaOSLab/EdgeNode/internal/utils/readers"
	"github.com/iwind/TeaGo/types"
	"io"
	"net/http"
	"net/url"
	"strings"
)

const (
	defaultM3u8EncryptParam = "ge_m3u8_token"
	m3u8EncryptExtension    = ".ge-m3u8-key"
)

// process hls
func (this *HTTPRequest) processHLSBefore() (blocked bool) {
	var requestPath = this.RawReq.URL.Path

	// .ge-m3u8-key
	if !strings.HasSuffix(requestPath, m3u8EncryptExtension) {
		return
	}

	blocked = true

	tokenBytes, base64Err := base64.StdEncoding.DecodeString(this.RawReq.URL.Query().Get(defaultM3u8EncryptParam))
	if base64Err != nil {
		this.writeCode(http.StatusBadRequest, "invalid m3u8 token: bad format", "invalid m3u8 token: bad format")
		return
	}

	if len(tokenBytes) != 32 {
		this.writeCode(http.StatusBadRequest, "invalid m3u8 token: bad length", "invalid m3u8 token: bad length")
		return
	}

	_, _ = this.writer.Write(tokenBytes[:16])

	return
}

// process m3u8 file
func (this *HTTPRequest) processM3u8Response(resp *http.Response) error {
	var requestPath = this.RawReq.URL.Path

	// .m3u8
	if strings.HasSuffix(requestPath, ".m3u8") &&
		this.web.HLS.Encrypting.MatchURL(this.URL()) {
		return this.processM3u8File(resp)
	}

	// .ts
	if strings.HasSuffix(requestPath, ".ts") {
		var token = this.RawReq.URL.Query().Get(defaultM3u8EncryptParam)
		if len(token) > 0 {
			return this.processTSFile(resp, token)
		}
	}

	return nil
}

func (this *HTTPRequest) processTSFile(resp *http.Response, token string) error {
	rawToken, err := base64.StdEncoding.DecodeString(token)
	if err != nil {
		return err
	}

	if len(rawToken) != 32 {
		return errors.New("invalid token length")
	}

	var key = rawToken[:16]
	var iv = rawToken[16:]
	block, err := aes.NewCipher(key)
	if err != nil {
		return fmt.Errorf("create cipher failed: %w", err)
	}
	var stream = cipher.NewCBCEncrypter(block, iv)
	var reader = readers.NewFilterReaderCloser(resp.Body)
	var blockSize = stream.BlockSize()
	reader.Add(func(p []byte, readErr error) error {
		var l = len(p)
		if l == 0 {
			return nil
		}

		var mod = l % blockSize
		if mod != 0 {
			p = append(p, bytes.Repeat([]byte{'0'}, blockSize-mod)...)
		}
		stream.CryptBlocks(p, p)
		return readErr
	})
	resp.Body = reader

	return nil
}

func (this *HTTPRequest) processM3u8File(resp *http.Response) error {
	const maxSize = 1 << 20

	// 检查内容长度
	if resp.Body == nil || resp.ContentLength == 0 || resp.ContentLength > maxSize {
		return nil
	}

	// 解压缩
	compressions.WrapHTTPResponse(resp)

	// 读取内容
	data, err := io.ReadAll(io.LimitReader(resp.Body, maxSize))
	if err != nil {
		_ = resp.Body.Close()
		return err
	}

	// 超出尺寸直接返回
	if len(data) == maxSize {
		resp.Body = io.NopCloser(io.MultiReader(bytes.NewBuffer(data), resp.Body))
		return nil
	}

	var lines = bytes.Split(data, []byte{'\n'})
	var addedKey = false

	var ivBytes = make([]byte, 16)
	var keyBytes = make([]byte, 16)

	_, ivErr := rand.Read(ivBytes)
	_, keyErr := rand.Read(keyBytes)
	if ivErr != nil || keyErr != nil {
		resp.Body = io.NopCloser(bytes.NewBuffer(data))
		return nil
	}

	var ivString = fmt.Sprintf("%x", ivBytes)
	var token = url.QueryEscape(base64.StdEncoding.EncodeToString(append(keyBytes, ivBytes...)))
	for index, line := range lines {
		if len(line) == 0 {
			continue
		}

		if bytes.HasPrefix(line, []byte("#EXT-X-KEY")) {
			goto returnData
		}

		if !addedKey && bytes.HasPrefix(line, []byte("#EXTINF")) {
			this.URL()
			var keyPath = strings.TrimSuffix(this.RawReq.URL.Path, ".m3u8") + m3u8EncryptExtension + "?" + defaultM3u8EncryptParam + "=" + token
			lines[index] = append([]byte("#EXT-X-KEY:METHOD=AES-128,URI=\""+this.requestScheme()+"://"+this.ReqHost+keyPath+"\",IV=0x"+ivString+
				"\n"), line...)
			addedKey = true
			continue
		}

		if line[0] != '#' && bytes.Contains(line, []byte(".ts")) {
			if bytes.ContainsRune(line, '?') {
				lines[index] = append(line, []byte("&"+defaultM3u8EncryptParam+"="+token)...)
			} else {
				lines[index] = append(line, []byte("?"+defaultM3u8EncryptParam+"="+token)...)
			}
		}
	}

	if addedKey {
		this.tags = append(this.tags, "m3u8")
	}

returnData:
	data = bytes.Join(lines, []byte{'\n'})
	resp.Body = io.NopCloser(bytes.NewBuffer(data))
	resp.ContentLength = int64(len(data))
	resp.Header.Set("Content-Length", types.String(resp.ContentLength))

	return nil
}