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

package nodes

import (
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/TeaOSLab/EdgeNode/internal/encryption"
	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
	"github.com/andybalholm/brotli"
)

const encryptionCacheVersion = "xor-v1"

// processPageEncryption 处理页面加密
func (this *HTTPRequest) processPageEncryption(resp *http.Response) error {
	// 首先检查是否是 /waf/loader.js，如果是则直接跳过（不应该被加密）
	// 这个检查必须在所有其他检查之前，确保 loader.js 永远不会被加密
	if strings.Contains(this.URL(), "/waf/loader.js") {
		return nil
	}

	if this.web.Encryption == nil {
		return nil
	}

	if !this.web.Encryption.IsOn {
		return nil
	}

	if !this.web.Encryption.IsEnabled() {
		return nil
	}

	// 检查 URL 白名单
	if this.web.Encryption.MatchExcludeURL(this.URL()) {
		return nil
	}

	// 检查 Content-Type 和 URL
	contentType := resp.Header.Get("Content-Type")
	contentTypeLower := strings.ToLower(contentType)
	urlLower := strings.ToLower(this.URL())

	var isHTML = strings.Contains(contentTypeLower, "text/html")
	// 判断是否为 JavaScript 文件：通过 Content-Type 或 URL 后缀
	var isJavaScript = strings.Contains(contentTypeLower, "text/javascript") ||
		strings.Contains(contentTypeLower, "application/javascript") ||
		strings.Contains(contentTypeLower, "application/x-javascript") ||
		strings.Contains(contentTypeLower, "text/ecmascript") ||
		strings.HasSuffix(urlLower, ".js") ||
		strings.Contains(urlLower, ".js?") ||
		strings.Contains(urlLower, ".js&")

	if !isHTML && !isJavaScript {
		return nil
	}

	// 检查内容大小（仅处理小于 10MB 的内容）
	if resp.ContentLength > 0 && resp.ContentLength > 10*1024*1024 {
		return nil
	}

	// 读取响应体
	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	_ = resp.Body.Close()

	// 如果源站返回了压缩内容，先解压再处理
	decodedBody, decoded, err := decodeResponseBody(bodyBytes, resp.Header.Get("Content-Encoding"))
	if err == nil && decoded {
		bodyBytes = decodedBody
		// 已经解压，移除 Content-Encoding
		resp.Header.Del("Content-Encoding")
	}

	// 检查实际大小
	if len(bodyBytes) > 10*1024*1024 {
		// 内容太大，恢复原始响应体
		resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
		return nil
	}

	var encryptedBytes []byte

	// 处理 JavaScript 文件
	if isJavaScript {
		// 检查是否需要加密独立的 JavaScript 文件
		if this.web.Encryption.Javascript == nil {
			resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			return nil
		}

		if !this.web.Encryption.Javascript.IsOn {
			resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			return nil
		}

		// 检查 URL 匹配
		if !this.web.Encryption.Javascript.MatchURL(this.URL()) {
			resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			return nil
		}

		// 跳过 Loader 文件
		if strings.Contains(this.URL(), "/waf/loader.js") ||
			strings.Contains(this.URL(), "waf-loader.js") ||
			strings.Contains(this.URL(), "__WAF_") {
			resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			return nil
		}

		encryptedBytes, err = this.encryptJavaScriptFile(bodyBytes, resp)
		if err != nil {
			remotelogs.Warn("HTTP_REQUEST_ENCRYPTION", "encrypt JavaScript file failed: "+err.Error())
			resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			return nil
		}
	} else if isHTML {
		// 加密 HTML 内容
		encryptedBytes, err = this.encryptHTMLScripts(bodyBytes, resp)
		if err != nil {
			remotelogs.Warn("HTTP_REQUEST_ENCRYPTION", "encrypt HTML failed: "+err.Error())
			// 加密失败，恢复原始响应体
			resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			return nil
		}
	} else {
		// 未知类型，恢复原始响应体
		resp.Body = io.NopCloser(bytes.NewReader(bodyBytes))
		return nil
	}

	// 替换响应体
	resp.Body = io.NopCloser(bytes.NewReader(encryptedBytes))
	resp.ContentLength = int64(len(encryptedBytes))
	resp.Header.Set("Content-Length", fmt.Sprintf("%d", len(encryptedBytes)))
	// 避免旧缓存导致解密算法不匹配
	resp.Header.Set("Cache-Control", "no-store, no-cache, must-revalidate")
	// 删除 Content-Encoding（如果存在），因为我们修改了内容
	resp.Header.Del("Content-Encoding")

	return nil
}

// encryptHTMLScripts 加密 HTML 中的脚本
func (this *HTTPRequest) encryptHTMLScripts(htmlBytes []byte, resp *http.Response) ([]byte, error) {
	html := string(htmlBytes)

	// 检查是否需要加密 HTML 脚本
	if this.web.Encryption.HTML == nil || !this.web.Encryption.HTML.IsOn {
		return htmlBytes, nil
	}

	// 检查 URL 匹配
	if !this.web.Encryption.HTML.MatchURL(this.URL()) {
		return htmlBytes, nil
	}

	// 生成密钥
	remoteIP := this.requestRemoteAddr(true)
	userAgent := this.RawReq.UserAgent()
	keyID, actualKey := encryption.GenerateEncryptionKey(remoteIP, userAgent, this.web.Encryption.KeyPolicy)

	// 检查缓存
	var cacheKey string
	if this.web.Encryption.Cache != nil && this.web.Encryption.Cache.IsOn {
		// 生成缓存键：keyID + URL + contentHash
		contentHash := fmt.Sprintf("%x", bytesHash(htmlBytes))
		cacheKey = fmt.Sprintf("encrypt_%s_%s_%s_%s", encryptionCacheVersion, keyID, this.URL(), contentHash)

		cache := encryption.SharedEncryptionCache(
			int(this.web.Encryption.Cache.MaxSize),
			time.Duration(this.web.Encryption.Cache.TTL)*time.Second,
		)

		if cached, ok := cache.Get(cacheKey); ok {
			return cached, nil
		}
	}

	// 提取并加密内联脚本
	if this.web.Encryption.HTML.EncryptInlineScripts {
		html = this.encryptInlineScripts(html, actualKey, keyID)
	}

	// 提取并加密外部脚本（通过 src 属性）
	if this.web.Encryption.HTML.EncryptExternalScripts {
		html = this.encryptExternalScripts(html, actualKey, keyID)
	}

	// 注入 Loader
	html = this.injectLoader(html)

	result := []byte(html)

	// 保存到缓存
	if this.web.Encryption.Cache != nil && this.web.Encryption.Cache.IsOn && len(cacheKey) > 0 {
		cache := encryption.SharedEncryptionCache(
			int(this.web.Encryption.Cache.MaxSize),
			time.Duration(this.web.Encryption.Cache.TTL)*time.Second,
		)
		cache.Set(cacheKey, result, this.web.Encryption.Cache.TTL)
	}

	return result, nil
}

// encryptInlineScripts 加密内联脚本
func (this *HTTPRequest) encryptInlineScripts(html string, key string, keyID string) string {
	// 匹配 <script>...</script>（不包含 src 属性）
	scriptRegex := regexp.MustCompile(`(?i)<script(?:\s+[^>]*)?>([\s\S]*?)</script>`)

	return scriptRegex.ReplaceAllStringFunc(html, func(match string) string {
		// 检查是否有 src 属性（外部脚本）
		if strings.Contains(strings.ToLower(match), "src=") {
			return match // 跳过外部脚本
		}

		// 提取脚本内容
		contentMatch := regexp.MustCompile(`(?i)<script(?:\s+[^>]*)?>([\s\S]*?)</script>`)
		matches := contentMatch.FindStringSubmatch(match)
		if len(matches) < 2 {
			return match
		}

		scriptContent := matches[1]

		// 跳过空脚本或仅包含空白字符的脚本
		if strings.TrimSpace(scriptContent) == "" {
			return match
		}

		// 跳过已加密的脚本（包含 __WAF_P__）
		if strings.Contains(scriptContent, "__WAF_P__") {
			return match
		}

		// 加密脚本内容
		encrypted, err := this.encryptScript(scriptContent, key)
		if err != nil {
			return match // 加密失败，返回原始内容
		}

		// 生成元数据（k 是 keyID，用于缓存；key 是实际密钥，用于解密）
		meta := fmt.Sprintf(`{"k":"%s","key":"%s","t":%d,"alg":"xor"}`, keyID, key, time.Now().Unix())

		// 替换为加密后的脚本（同步解密执行，保证脚本顺序）
		return fmt.Sprintf(
			`<script>(function(){
				function xorDecodeToString(b64,key){
					var bin=atob(b64);
					var out=new Uint8Array(bin.length);
					for(var i=0;i<bin.length;i++){out[i]=bin.charCodeAt(i)^key.charCodeAt(i%%key.length);}
					if (typeof TextDecoder !== 'undefined') {
						return new TextDecoder().decode(out);
					}
					var s='';for(var j=0;j<out.length;j++){s+=String.fromCharCode(out[j]);}
					return s;
				}
				try{var meta=%s;var code=xorDecodeToString("%s",meta.key);window.eval(code);}catch(e){console.error('WAF inline decrypt/execute failed',e);}
			})();</script>`,
			meta,
			encrypted,
		)
	})
}

// encryptExternalScripts 加密外部脚本（通过替换 src 为加密后的内容）
// 注意：这里我们实际上是将外部脚本的内容内联并加密
func (this *HTTPRequest) encryptExternalScripts(html string, key string, keyID string) string {
	// 匹配 <script src="..."></script>
	scriptRegex := regexp.MustCompile(`(?i)<script\s+([^>]*src\s*=\s*["']([^"']+)["'][^>]*)>\s*</script>`)

	return scriptRegex.ReplaceAllStringFunc(html, func(match string) string {
		// 提取 src URL
		srcMatch := regexp.MustCompile(`(?i)src\s*=\s*["']([^"']+)["']`)
		srcMatches := srcMatch.FindStringSubmatch(match)
		if len(srcMatches) < 2 {
			return match
		}

		srcURL := srcMatches[1]

		// 跳过已加密的脚本或 Loader
		if strings.Contains(srcURL, "waf-loader.js") || strings.Contains(srcURL, "__WAF_") {
			return match
		}

		// 注意：实际实现中，我们需要获取外部脚本的内容
		// 这里为了简化，我们只是标记需要加密，实际内容获取需要在响应处理时进行
		// 当前实现：将外部脚本转换为内联加密脚本的占位符
		// 实际生产环境需要：1. 获取外部脚本内容 2. 加密 3. 替换

		return match // 暂时返回原始内容，后续可以扩展
	})
}

// encryptScript 加密脚本内容
func (this *HTTPRequest) encryptScript(scriptContent string, key string) (string, error) {
	// 1. XOR 加密（不压缩，避免浏览器解压依赖导致顺序问题）
	encrypted := xorEncrypt([]byte(scriptContent), []byte(key))

	// 2. Base64 编码
	encoded := base64.StdEncoding.EncodeToString(encrypted)

	return encoded, nil
}

// injectLoader 注入 Loader 脚本
func (this *HTTPRequest) injectLoader(html string) string {
	// 检查是否已经注入
	if strings.Contains(html, "waf-loader.js") {
		return html
	}

	// 在 </head> 之前注入，如果没有 </head>，则在 </body> 之前注入
	// 不使用 async，确保 loader 在解析阶段优先加载并执行
	loaderScript := `<script src="/waf/loader.js"></script>`

	if strings.Contains(html, "</head>") {
		return strings.Replace(html, "</head>", loaderScript+"</head>", 1)
	} else if strings.Contains(html, "</body>") {
		return strings.Replace(html, "</body>", loaderScript+"</body>", 1)
	} else {
		// 如果都没有，在开头注入
		return loaderScript + html
	}
}

// compressGzip 使用 Gzip 压缩（浏览器原生支持）
func compressGzip(data []byte) ([]byte, error) {
	var buf bytes.Buffer
	writer := gzip.NewWriter(&buf)
	_, err := writer.Write(data)
	if err != nil {
		writer.Close()
		return nil, err
	}
	err = writer.Close()
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// decodeResponseBody 根据 Content-Encoding 解压响应体
func decodeResponseBody(body []byte, encoding string) ([]byte, bool, error) {
	enc := strings.ToLower(strings.TrimSpace(encoding))
	if enc == "" || enc == "identity" {
		return body, false, nil
	}
	switch enc {
	case "gzip":
		reader, err := gzip.NewReader(bytes.NewReader(body))
		if err != nil {
			return body, false, err
		}
		defer reader.Close()
		decoded, err := io.ReadAll(reader)
		if err != nil {
			return body, false, err
		}
		return decoded, true, nil
	case "br":
		reader := brotli.NewReader(bytes.NewReader(body))
		decoded, err := io.ReadAll(reader)
		if err != nil {
			return body, false, err
		}
		return decoded, true, nil
	default:
		// 未知编码，保持原样
		return body, false, nil
	}
}

// compressBrotli 使用 Brotli 压缩（保留用于其他用途）
func compressBrotli(data []byte, level int) ([]byte, error) {
	var buf bytes.Buffer
	writer := brotli.NewWriterOptions(&buf, brotli.WriterOptions{
		Quality: level,
		LGWin:   14,
	})
	_, err := writer.Write(data)
	if err != nil {
		writer.Close()
		return nil, err
	}
	err = writer.Close()
	if err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// xorEncrypt XOR 加密
func xorEncrypt(data []byte, key []byte) []byte {
	result := make([]byte, len(data))
	keyLen := len(key)
	if keyLen == 0 {
		return data
	}

	for i := 0; i < len(data); i++ {
		result[i] = data[i] ^ key[i%keyLen]
	}

	return result
}

// encryptJavaScriptFile 加密独立的 JavaScript 文件
func (this *HTTPRequest) encryptJavaScriptFile(jsBytes []byte, resp *http.Response) ([]byte, error) {
	jsContent := string(jsBytes)

	// 跳过空文件
	if strings.TrimSpace(jsContent) == "" {
		return jsBytes, nil
	}

	// 跳过已加密的脚本（包含 __WAF_P__）
	if strings.Contains(jsContent, "__WAF_P__") {
		return jsBytes, nil
	}

	// 生成密钥
	remoteIP := this.requestRemoteAddr(true)
	userAgent := this.RawReq.UserAgent()
	keyID, actualKey := encryption.GenerateEncryptionKey(remoteIP, userAgent, this.web.Encryption.KeyPolicy)

	// 检查缓存
	var cacheKey string
	if this.web.Encryption.Cache != nil && this.web.Encryption.Cache.IsOn {
		// 生成缓存键：keyID + URL + contentHash
		contentHash := fmt.Sprintf("%x", bytesHash(jsBytes))
		cacheKey = fmt.Sprintf("encrypt_js_%s_%s_%s_%s", encryptionCacheVersion, keyID, this.URL(), contentHash)

		cache := encryption.SharedEncryptionCache(
			int(this.web.Encryption.Cache.MaxSize),
			time.Duration(this.web.Encryption.Cache.TTL)*time.Second,
		)

		if cached, ok := cache.Get(cacheKey); ok {
			return cached, nil
		}
	}

	// 加密脚本内容
	encrypted, err := this.encryptScript(jsContent, actualKey)
	if err != nil {
		return nil, err
	}

	// 生成元数据（k 是 keyID，用于缓存；key 是实际密钥，用于解密）
	meta := fmt.Sprintf(`{"k":"%s","key":"%s","t":%d,"alg":"xor"}`, keyID, actualKey, time.Now().Unix())

	// 生成加密后的 JavaScript 代码（同步解密执行，保证脚本顺序）
	encryptedJS := fmt.Sprintf(`(function() {
		try {
			function xorDecodeToString(b64, key) {
				var bin = atob(b64);
				var out = new Uint8Array(bin.length);
				for (var i = 0; i < bin.length; i++) {
					out[i] = bin.charCodeAt(i) ^ key.charCodeAt(i %% key.length);
				}
				if (typeof TextDecoder !== 'undefined') {
					return new TextDecoder().decode(out);
				}
				var s = '';
				for (var j = 0; j < out.length; j++) {
					s += String.fromCharCode(out[j]);
				}
				return s;
			}
			var meta = %s;
			var code = xorDecodeToString("%s", meta.key);
			// 使用全局 eval，尽量保持和 <script> 一致的作用域
			window.eval(code);
		} catch (e) {
			console.error('WAF JS decrypt/execute failed', e);
		}
	})();`, meta, encrypted)

	result := []byte(encryptedJS)

	// 保存到缓存
	if this.web.Encryption.Cache != nil && this.web.Encryption.Cache.IsOn && len(cacheKey) > 0 {
		cache := encryption.SharedEncryptionCache(
			int(this.web.Encryption.Cache.MaxSize),
			time.Duration(this.web.Encryption.Cache.TTL)*time.Second,
		)
		cache.Set(cacheKey, result, this.web.Encryption.Cache.TTL)
	}

	return result, nil
}

// bytesHash 计算字节数组的简单哈希（用于缓存键）
func bytesHash(data []byte) uint64 {
	var hash uint64 = 5381
	for _, b := range data {
		hash = ((hash << 5) + hash) + uint64(b)
	}
	return hash
}