// Package rules implements match-and-replace: user-defined rules that
// rewrite request/response headers or bodies as they pass through the
// proxy. See ApplyHeaders and ApplyBody for how each part is handled.
package rules

import (
	"bufio"
	"net/http"
	"net/textproto"
	"regexp"
	"sort"
	"strings"
)

// Rule is one match-and-replace rule.
type Rule struct {
	ID      int64  `json:"id"`
	Enabled bool   `json:"enabled"`
	Name    string `json:"name"`
	Scope   string `json:"scope"` // "request" or "response"
	Part    string `json:"part"`  // "header" or "body"
	Match   string `json:"match"`
	Replace string `json:"replace"`
	IsRegex bool   `json:"is_regex"`
	// Position orders rule application (ascending) when several rules
	// could touch the same text.
	Position int `json:"position"`
}

// ApplyHeaders rewrites h in place by serializing it to a raw
// "Name: value\r\n" block, running every enabled rule with Part=="header"
// over that text (in Position order), and reparsing the result. Working
// on the raw text rather than per-value substitution is what lets a rule
// add or remove a header entirely, not just rewrite an existing value -
// matching how Burp's header match/replace works. If a rule's output
// doesn't parse back as valid headers, ApplyHeaders returns h unchanged
// rather than risk sending something corrupted.
func ApplyHeaders(h http.Header, rs []Rule) http.Header {
	keys := make([]string, 0, len(h))
	for k := range h {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var block strings.Builder
	for _, k := range keys {
		for _, v := range h[k] {
			block.WriteString(k)
			block.WriteString(": ")
			block.WriteString(v)
			block.WriteString("\r\n")
		}
	}
	text := block.String()

	changed := false
	for _, r := range sortedByPosition(rs) {
		if !r.Enabled || r.Part != "header" {
			continue
		}
		if next, ok := apply(text, r); ok {
			text, changed = next, true
		}
	}
	if !changed {
		return h
	}

	tp := textproto.NewReader(bufio.NewReader(strings.NewReader(text + "\r\n")))
	mh, err := tp.ReadMIMEHeader()
	if err != nil {
		return h
	}
	return http.Header(mh)
}

// HasBodyRules reports whether any of rs (already filtered to enabled
// rules for one scope, as returned by store.EnabledRules) is a body
// rule - used to decide whether materializing a request/response body
// at all is worth it. Most traffic has no body rule configured and
// should keep streaming through unbuffered, which is what makes exact
// capture of arbitrarily large bodies possible in the first place.
func HasBodyRules(rs []Rule) bool {
	for _, r := range rs {
		if r.Enabled && r.Part == "body" {
			return true
		}
	}
	return false
}

// ApplyBody rewrites body by running every enabled rule with
// Part=="body" over it, in Position order - literal substring replace
// by default, or regex if the rule's Regex toggle is on. Unlike
// ApplyHeaders, this operates directly on raw bytes with no reparse/
// validation step: body content isn't a fixed format the way a header
// block is, so there's nothing to parse back and confirm is still
// well-formed. A rule that doesn't match anything is simply a no-op,
// same as ApplyHeaders.
func ApplyBody(body []byte, rs []Rule) []byte {
	text := string(body)
	for _, r := range sortedByPosition(rs) {
		if !r.Enabled || r.Part != "body" {
			continue
		}
		if next, ok := apply(text, r); ok {
			text = next
		}
	}
	return []byte(text)
}

func sortedByPosition(rs []Rule) []Rule {
	out := make([]Rule, len(rs))
	copy(out, rs)
	sort.SliceStable(out, func(i, j int) bool { return out[i].Position < out[j].Position })
	return out
}

func apply(text string, r Rule) (string, bool) {
	if r.IsRegex {
		re, err := regexp.Compile(r.Match)
		if err != nil {
			return text, false
		}
		return re.ReplaceAllString(text, r.Replace), true
	}
	if r.Match == "" {
		return text, false
	}
	return strings.ReplaceAll(text, r.Match, r.Replace), true
}