//proxy/rewriter.go
package proxy

import (
    "bytes"
    "golang.org/x/net/html"
    "net/url"
    "strings"
)

type ContentRewriter struct {
    baseURL *url.URL
}

func NewContentRewriter(baseURL string) (*ContentRewriter, error) {
    u, err := url.Parse(baseURL)
    if err != nil {
        return nil, err
    }
    
    return &ContentRewriter{
        baseURL: u,
    }, nil
}

// RewriteHTML 重写 HTML 内容中的所有 URL
func (r *ContentRewriter) RewriteHTML(body []byte) ([]byte, error) {
    doc, err := html.Parse(bytes.NewReader(body))
    if err != nil {
        // 如果解析失败，使用简单的字符串替换
        return r.simpleRewriteHTML(body), nil
    }
    
    r.rewriteNode(doc)
    
    var buf bytes.Buffer
    if err := html.Render(&buf, doc); err != nil {
        return r.simpleRewriteHTML(body), nil
    }
    
    return buf.Bytes(), nil
}

// rewriteNode 递归重写 HTML 节点
func (r *ContentRewriter) rewriteNode(n *html.Node) {
    if n.Type == html.ElementNode {
        // 重写需要处理的属性
        attrs := map[string]bool{
            "href":   true,
            "src":    true,
            "action": true,
            "data":   true,
        }
        
        for i, attr := range n.Attr {
            if attrs[attr.Key] {
                if rewritten := r.rewriteURL(attr.Val); rewritten != attr.Val {
                    n.Attr[i].Val = rewritten
                }
            }
            
            // 处理 srcset 属性
            if attr.Key == "srcset" {
                n.Attr[i].Val = r.rewriteSrcset(attr.Val)
            }
            
            // 处理 style 属性中的 URL
            if attr.Key == "style" {
                n.Attr[i].Val = r.rewriteInlineCSS(attr.Val)
            }
        }
        
        // 处理 <base> 标签
        if n.Data == "base" {
            for i, attr := range n.Attr {
                if attr.Key == "href" {
                    n.Attr[i].Val = r.baseURL.String()
                }
            }
        }
        
        // 处理 <style> 标签内容
        if n.Data == "style" && n.FirstChild != nil {
            if n.FirstChild.Type == html.TextNode {
                n.FirstChild.Data = r.rewriteInlineCSS(n.FirstChild.Data)
            }
        }
        
        // 处理 <script> 标签，移除可能的跟踪脚本
        if n.Data == "script" {
            for _, attr := range n.Attr {
                if attr.Key == "src" {
                    // 可以在这里过滤掉已知的跟踪脚本
                    if r.isTrackingScript(attr.Val) {
                        // 移除此节点
                        if n.Parent != nil {
                            n.Parent.RemoveChild(n)
                            return
                        }
                    }
                }
            }
        }
    }
    
    // 递归处理子节点
    for c := n.FirstChild; c != nil; c = c.NextSibling {
        r.rewriteNode(c)
    }
}

// rewriteURL 重写单个 URL
func (r *ContentRewriter) rewriteURL(urlStr string) string {
    urlStr = strings.TrimSpace(urlStr)
    
    // 跳过特殊协议
    if strings.HasPrefix(urlStr, "javascript:") ||
        strings.HasPrefix(urlStr, "data:") ||
        strings.HasPrefix(urlStr, "mailto:") ||
        strings.HasPrefix(urlStr, "tel:") ||
        strings.HasPrefix(urlStr, "#") ||
        urlStr == "" {
        return urlStr
    }
    
    // 解析 URL
    u, err := url.Parse(urlStr)
    if err != nil {
        return urlStr
    }
    
    // 如果是相对 URL，转换为绝对 URL
    if !u.IsAbs() {
        u = r.baseURL.ResolveReference(u)
    }
    
    // 生成代理 URL
    return "/proxy?url=" + url.QueryEscape(u.String())
}

// rewriteSrcset 重写 srcset 属性
func (r *ContentRewriter) rewriteSrcset(srcset string) string {
    if srcset == "" {
        return srcset
    }
    
    parts := strings.Split(srcset, ",")
    var rewritten []string
    
    for _, part := range parts {
        part = strings.TrimSpace(part)
        fields := strings.Fields(part)
        
        if len(fields) > 0 {
            fields[0] = r.rewriteURL(fields[0])
            rewritten = append(rewritten, strings.Join(fields, " "))
        }
    }
    
    return strings.Join(rewritten, ", ")
}

// RewriteCSS 重写 CSS 内容
func (r *ContentRewriter) RewriteCSS(body []byte) []byte {
    content := string(body)
    return []byte(r.rewriteInlineCSS(content))
}

// rewriteInlineCSS 重写内联 CSS 中的 URL
func (r *ContentRewriter) rewriteInlineCSS(css string) string {
    // 匹配 url(...) 模式
    result := css
    
    // 处理 url("...") 和 url('...') 和 url(...)
    patterns := []struct {
        prefix string
        suffix string
    }{
        {`url("`, `")`},
        {`url('`, `')`},
        {`url(`, `)`},
    }
    
    for _, pattern := range patterns {
        start := 0
        for {
            idx := strings.Index(result[start:], pattern.prefix)
            if idx == -1 {
                break
            }
            
            idx += start
            urlStart := idx + len(pattern.prefix)
            urlEnd := strings.Index(result[urlStart:], pattern.suffix)
            
            if urlEnd == -1 {
                break
            }
            
            urlEnd += urlStart
            originalURL := result[urlStart:urlEnd]
            rewrittenURL := r.rewriteURL(originalURL)
            
            result = result[:urlStart] + rewrittenURL + result[urlEnd:]
            start = urlStart + len(rewrittenURL)
        }
    }
    
    // 处理 @import
    result = r.rewriteImports(result)
    
    return result
}

// rewriteImports 重写 CSS @import 语句
func (r *ContentRewriter) rewriteImports(css string) string {
    result := css
    
    patterns := []string{
        `@import "`,
        `@import '`,
        `@import url("`,
        `@import url('`,
    }
    
    for _, pattern := range patterns {
        start := 0
        for {
            idx := strings.Index(result[start:], pattern)
            if idx == -1 {
                break
            }
            
            idx += start
            urlStart := idx + len(pattern)
            
            var endChar string
            if strings.Contains(pattern, `"`) {
                endChar = `"`
            } else {
                endChar = `'`
            }
            
            urlEnd := strings.Index(result[urlStart:], endChar)
            if urlEnd == -1 {
                break
            }
            
            urlEnd += urlStart
            originalURL := result[urlStart:urlEnd]
            rewrittenURL := r.rewriteURL(originalURL)
            
            result = result[:urlStart] + rewrittenURL + result[urlEnd:]
            start = urlStart + len(rewrittenURL)
        }
    }
    
    return result
}

// simpleRewriteHTML 简单的字符串替换重写（备用方案）
func (r *ContentRewriter) simpleRewriteHTML(body []byte) []byte {
    content := string(body)
    
    // 重写绝对 URL
    baseStr := r.baseURL.Scheme + "://" + r.baseURL.Host
    
    replacements := []struct {
        old string
        new string
    }{
        {`href="` + baseStr, `href="/proxy?url=` + url.QueryEscape(baseStr)},
        {`src="` + baseStr, `src="/proxy?url=` + url.QueryEscape(baseStr)},
        {`action="` + baseStr, `action="/proxy?url=` + url.QueryEscape(baseStr)},
        {`href='` + baseStr, `href='/proxy?url=` + url.QueryEscape(baseStr)},
        {`src='` + baseStr, `src='/proxy?url=` + url.QueryEscape(baseStr)},
    }
    
    for _, r := range replacements {
        content = strings.ReplaceAll(content, r.old, r.new)
    }
    
    return []byte(content)
}

// isTrackingScript 检查是否是跟踪脚本
func (r *ContentRewriter) isTrackingScript(src string) bool {
    trackingDomains := []string{
        "google-analytics.com",
        "googletagmanager.com",
        "facebook.net",
        "doubleclick.net",
        "analytics.js",
        "ga.js",
        "gtag.js",
    }
    
    srcLower := strings.ToLower(src)
    for _, domain := range trackingDomains {
        if strings.Contains(srcLower, domain) {
            return true
        }
    }
    
    return false
}