package rewriter

import (
	"net/http"
	"regexp"
	"strings"
)

// Rewriter URL重写器
// 用于在反向代理中重写请求URL
type Rewriter struct {
	// 重写规则列表
	rules []*RewriteRule
}

// RewriteRule 重写规则
type RewriteRule struct {
	// 匹配模式
	Pattern string
	// 替换模式
	Replacement string
	// 是否使用正则表达式
	UseRegex bool
	// 编译后的正则表达式
	regex *regexp.Regexp
}

// NewRewriter 创建URL重写器
func NewRewriter() *Rewriter {
	return &Rewriter{
		rules: make([]*RewriteRule, 0),
	}
}

// AddRule 添加重写规则
func (r *Rewriter) AddRule(pattern, replacement string, useRegex bool) error {
	rule := &RewriteRule{
		Pattern:     pattern,
		Replacement: replacement,
		UseRegex:    useRegex,
	}

	if useRegex {
		regex, err := regexp.Compile(pattern)
		if err != nil {
			return err
		}
		rule.regex = regex
	}

	r.rules = append(r.rules, rule)
	return nil
}

// Rewrite 重写URL
func (r *Rewriter) Rewrite(req *http.Request) {
	path := req.URL.Path

	for _, rule := range r.rules {
		if rule.UseRegex {
			if rule.regex.MatchString(path) {
				req.URL.Path = rule.regex.ReplaceAllString(path, rule.Replacement)
				break
			}
		} else {
			if strings.HasPrefix(path, rule.Pattern) {
				req.URL.Path = strings.Replace(path, rule.Pattern, rule.Replacement, 1)
				break
			}
		}
	}
}

// RewriteResponse 重写响应
// 主要用于处理响应中的Location头和内容中的URL
func (r *Rewriter) RewriteResponse(resp *http.Response, originHost string) {
	// 处理重定向头
	location := resp.Header.Get("Location")
	if location != "" {
		// 将后端服务器的域名替换成代理服务器的域名
		for _, rule := range r.rules {
			if rule.UseRegex && rule.regex != nil {
				if rule.regex.MatchString(location) {
					newLocation := rule.regex.ReplaceAllString(location, rule.Replacement)
					resp.Header.Set("Location", newLocation)
					break
				}
			}
		}
	}
}

// LoadRulesFromFile 从文件加载重写规则
func (r *Rewriter) LoadRulesFromFile(filename string) error {
	// 实现从配置文件加载规则的逻辑
	// 这里省略实现细节
	return nil
}