package forwardplugin

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
)

func init() {
	Register("cookie_domain", newCookieDomainPlugin)
}

type cookieDomainConfig struct {
	From string `json:"from"`
	To   string `json:"to"`
}

type cookieDomainPlugin struct {
	from string
	to   string
}

func newCookieDomainPlugin(raw json.RawMessage) (Plugin, error) {
	var cfg cookieDomainConfig
	if err := json.Unmarshal(raw, &cfg); err != nil {
		return nil, fmt.Errorf("cookie_domain 配置无效")
	}
	return &cookieDomainPlugin{from: cfg.From, to: cfg.To}, nil
}

func (p *cookieDomainPlugin) Type() string { return "cookie_domain" }

func (p *cookieDomainPlugin) ValidateConfig(raw json.RawMessage) error {
	var cfg cookieDomainConfig
	if err := json.Unmarshal(raw, &cfg); err != nil {
		return fmt.Errorf("cookie_domain 配置无效")
	}
	if strings.TrimSpace(cfg.From) == "" {
		return fmt.Errorf("cookie_domain.from 不能为空")
	}
	if strings.TrimSpace(cfg.To) == "" {
		return fmt.Errorf("cookie_domain.to 不能为空")
	}
	return nil
}

func (p *cookieDomainPlugin) NeedsBody() bool { return false }

func (p *cookieDomainPlugin) HandleRequest(_ *Context, _ *http.Request) (bool, error) {
	return false, nil
}

func (p *cookieDomainPlugin) HandleResponse(_ *Context, resp *BufferedResponse) error {
	cookies := resp.Header.Values("Set-Cookie")
	if len(cookies) == 0 {
		return nil
	}
	resp.Header.Del("Set-Cookie")
	for _, c := range cookies {
		resp.Header.Add("Set-Cookie", rewriteCookieDomain(c, p.from, p.to))
	}
	return nil
}

func rewriteCookieDomain(cookie, from, to string) string {
	parts := strings.Split(cookie, ";")
	for i, part := range parts {
		trimmed := strings.TrimSpace(part)
		kv := strings.SplitN(trimmed, "=", 2)
		if len(kv) != 2 {
			continue
		}
		if strings.EqualFold(kv[0], "domain") {
			val := strings.TrimSpace(kv[1])
			if val == from || strings.TrimPrefix(val, ".") == strings.TrimPrefix(from, ".") {
				parts[i] = "Domain=" + to
			} else {
				parts[i] = trimmed
			}
		} else {
			parts[i] = trimmed
		}
	}
	return strings.Join(parts, "; ")
}