package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"sort"
	"strings"
	"time"

	"gopkg.in/yaml.v3"
)

const managedPrefix = "3x-rule-sync:"

type Config struct {
	Servers []Server `yaml:"servers"`
	Rules   []Rule   `yaml:"rules"`
	Sync    Sync     `yaml:"sync"`
}

type Server struct {
	Name     string `yaml:"name"`
	URL      string `yaml:"url"`
	Token    string `yaml:"token"`
	TokenEnv string `yaml:"tokenEnv"`
}

type Rule struct {
	Name     string         `yaml:"name"`
	Disabled bool           `yaml:"disabled"`
	Match    map[string]any `yaml:"match"`
	Outbound Outbound       `yaml:"outbound"`
}

type Outbound struct {
	TagRegex string `yaml:"tagRegex"`
	Strategy string `yaml:"strategy"`
}

type Sync struct {
	Interval string `yaml:"interval"`
	Position string `yaml:"position"`
}

type apiResponse struct {
	Success bool            `json:"success"`
	Msg     string          `json:"msg"`
	Obj     json.RawMessage `json:"obj"`
}

type xrayTemplate struct {
	config                   map[string]any
	subscriptionOutboundTags []string
	outboundTestURL          string
}

func main() {
	configPath := flag.String("config", "config.yaml", "path to YAML configuration")
	once := flag.Bool("once", false, "run one synchronization and exit")
	dryRun := flag.Bool("dry-run", false, "show changes without saving")
	debug := flag.Bool("debug", false, "log available outbound tags")
	flag.Parse()

	cfg, err := loadConfig(*configPath)
	if err != nil {
		log.Fatal(err)
	}
	if err := syncAll(context.Background(), cfg, *dryRun, *debug); err != nil {
		log.Fatal(err)
	}
	if *once || cfg.Sync.Interval == "" {
		return
	}
	interval, err := time.ParseDuration(cfg.Sync.Interval)
	if err != nil || interval <= 0 {
		log.Fatalf("sync.interval must be a positive duration: %q", cfg.Sync.Interval)
	}
	for range time.Tick(interval) {
		if err := syncAll(context.Background(), cfg, *dryRun, *debug); err != nil {
			log.Printf("sync failed: %v", err)
		}
	}
}

func loadConfig(path string) (Config, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return Config{}, err
	}
	var cfg Config
	decoder := yaml.NewDecoder(bytes.NewReader(b))
	decoder.KnownFields(true)
	if err := decoder.Decode(&cfg); err != nil {
		return Config{}, fmt.Errorf("parse config: %w", err)
	}
	if len(cfg.Servers) == 0 {
		return Config{}, errors.New("servers must not be empty")
	}
	if cfg.Sync.Position == "" {
		cfg.Sync.Position = "last"
	}
	if cfg.Sync.Position != "first" && cfg.Sync.Position != "last" {
		return Config{}, errors.New("sync.position must be first or last")
	}
	seen := map[string]bool{}
	for _, rule := range cfg.Rules {
		if rule.Name == "" {
			return Config{}, errors.New("each rule needs name")
		}
		if seen[rule.Name] {
			return Config{}, fmt.Errorf("duplicate rule name %q", rule.Name)
		}
		seen[rule.Name] = true
		if rule.Outbound.TagRegex == "" {
			return Config{}, fmt.Errorf("rule %q needs outbound.tagRegex", rule.Name)
		}
		if _, err := regexp.Compile(rule.Outbound.TagRegex); err != nil {
			return Config{}, fmt.Errorf("rule %q: invalid tagRegex: %w", rule.Name, err)
		}
		if rule.Outbound.Strategy == "" {
			continue
		}
		if !validStrategy(rule.Outbound.Strategy) {
			return Config{}, fmt.Errorf("rule %q: unsupported strategy %q", rule.Name, rule.Outbound.Strategy)
		}
	}
	for _, server := range cfg.Servers {
		if server.Name == "" || server.URL == "" {
			return Config{}, errors.New("each server needs name and url")
		}
		if server.Token == "" && server.TokenEnv == "" {
			return Config{}, fmt.Errorf("server %q needs token or tokenEnv", server.Name)
		}
	}
	return cfg, nil
}

func syncAll(ctx context.Context, cfg Config, dryRun, debug bool) error {
	var failures []string
	for _, server := range cfg.Servers {
		if err := syncServer(ctx, server, cfg.Rules, cfg.Sync.Position, dryRun, debug); err != nil {
			failures = append(failures, fmt.Sprintf("%s: %v", server.Name, err))
		} else {
			log.Printf("%s: synchronized", server.Name)
		}
	}
	if len(failures) > 0 {
		return errors.New(strings.Join(failures, "; "))
	}
	return nil
}

func syncServer(ctx context.Context, server Server, rules []Rule, position string, dryRun, debug bool) error {
	token := server.Token
	if server.TokenEnv != "" && os.Getenv(server.TokenEnv) != "" {
		token = os.Getenv(server.TokenEnv)
	}
	if token == "" {
		return fmt.Errorf("token is empty")
	}
	client := &panelClient{baseURL: strings.TrimRight(server.URL, "/"), token: token, http: &http.Client{Timeout: 30 * time.Second}}
	template, err := client.getTemplate(ctx)
	if err != nil {
		return err
	}
	config := template.config
	available, err := tags(config["outbounds"])
	if err != nil {
		return err
	}
	available = uniqueTags(append(available, template.subscriptionOutboundTags...))
	if len(available) == 0 {
		return errors.New("runtime Xray JSON contains no outbounds")
	}
	if debug {
		log.Printf("%s: available outbounds: %s", server.Name, strings.Join(available, ", "))
	}
	updated, err := reconcileWithOutbounds(config, rules, position, available)
	if err != nil {
		return err
	}
	if reflect.DeepEqual(config, updated) {
		log.Printf("%s: already up to date", server.Name)
		return nil
	}
	if dryRun {
		log.Printf("%s: would update Xray configuration", server.Name)
		return nil
	}
	return client.updateConfig(ctx, updated, template.outboundTestURL)
}

type panelClient struct {
	baseURL, token string
	http           *http.Client
}

func (c *panelClient) getTemplate(ctx context.Context) (xrayTemplate, error) {
	var response apiResponse
	if err := c.post(ctx, "/panel/api/xray/", nil, &response); err != nil {
		return xrayTemplate{}, err
	}
	if !response.Success {
		return xrayTemplate{}, fmt.Errorf("panel API: %s", response.Msg)
	}
	return parseTemplate(response.Obj)
}

func parseTemplate(raw json.RawMessage) (xrayTemplate, error) {
	var text string
	if json.Unmarshal(raw, &text) == nil {
		raw = json.RawMessage(text)
	}
	var envelope map[string]json.RawMessage
	if err := json.Unmarshal(raw, &envelope); err != nil {
		return xrayTemplate{}, fmt.Errorf("decode Xray settings response: %w", err)
	}
	setting, wrapped := envelope["xraySetting"]
	if !wrapped {
		setting = raw
	}
	config, err := decodeConfig(setting)
	if err != nil {
		return xrayTemplate{}, err
	}
	result := xrayTemplate{config: config}
	if value, ok := envelope["subscriptionOutboundTags"]; ok {
		_ = json.Unmarshal(value, &result.subscriptionOutboundTags)
	}
	if value, ok := envelope["outboundTestUrl"]; ok {
		_ = json.Unmarshal(value, &result.outboundTestURL)
	}
	return result, nil
}

func decodeConfig(raw json.RawMessage) (map[string]any, error) {
	var text string
	if json.Unmarshal(raw, &text) == nil {
		raw = json.RawMessage(text)
	}
	var config map[string]any
	if err := json.Unmarshal(raw, &config); err != nil {
		return nil, fmt.Errorf("decode Xray JSON: %w", err)
	}
	return config, nil
}

func (c *panelClient) updateConfig(ctx context.Context, config map[string]any, outboundTestURL string) error {
	b, err := json.Marshal(config)
	if err != nil {
		return err
	}
	var response apiResponse
	err = c.post(ctx, "/panel/api/xray/update", url.Values{"xraySetting": {string(b)}, "outboundTestUrl": {outboundTestURL}}, &response)
	if errors.Is(err, io.EOF) {
		return nil
	}
	if err != nil {
		return err
	}
	if !response.Success {
		return fmt.Errorf("panel API: %s", response.Msg)
	}
	return nil
}

func (c *panelClient) post(ctx context.Context, path string, form url.Values, out any) error {
	var body io.Reader
	if form != nil {
		body = strings.NewReader(form.Encode())
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+path, body)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+c.token)
	if form != nil {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}
	resp, err := c.http.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		return fmt.Errorf("HTTP %s", resp.Status)
	}
	if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
		return fmt.Errorf("POST %s: decode response: %w", path, err)
	}
	return nil
}

func reconcile(config map[string]any, desired []Rule, position string) (map[string]any, error) {
	outboundTags, err := tags(config["outbounds"])
	if err != nil {
		return nil, err
	}
	return reconcileWithOutbounds(config, desired, position, outboundTags)
}

func reconcileWithOutbounds(config map[string]any, desired []Rule, position string, outboundTags []string) (map[string]any, error) {
	updated := clone(config)
	routing := object(updated, "routing")
	existingRules, err := objects(routing["rules"])
	if err != nil {
		return nil, err
	}
	manualRules := filterUnmanaged(existingRules, "ruleTag")
	existingBalancers, err := objects(routing["balancers"])
	if err != nil {
		return nil, err
	}
	manualBalancers := filterUnmanaged(existingBalancers, "tag")

	managedRules := make([]map[string]any, 0, len(desired))
	managedBalancers := []map[string]any{}
	for _, rule := range desired {
		matches := matchingTags(outboundTags, rule.Outbound.TagRegex)
		if len(matches) == 0 {
			return nil, fmt.Errorf("rule %q: no outbound tag matches %q", rule.Name, rule.Outbound.TagRegex)
		}
		item := clone(rule.Match)
		delete(item, "outboundTag")
		delete(item, "balancerTag")
		delete(item, "ruleTag")
		item["ruleTag"] = managedPrefix + rule.Name
		item["enabled"] = !rule.Disabled
		if len(matches) == 1 {
			item["outboundTag"] = matches[0]
		} else {
			balancerTag := managedPrefix + "balancer:" + rule.Name
			item["balancerTag"] = balancerTag
			strategy := rule.Outbound.Strategy
			if strategy == "" {
				strategy = "leastPing"
			}
			if (strategy == "leastPing" || strategy == "leastLoad") && updated["observatory"] == nil {
				return nil, errors.New("leastPing and leastLoad require an Xray observatory")
			}
			managedBalancers = append(managedBalancers, map[string]any{"tag": balancerTag, "selector": matches, "strategy": map[string]any{"type": strategy}})
		}
		managedRules = append(managedRules, item)
	}
	if position == "first" {
		routing["rules"] = append(managedRules, manualRules...)
	} else {
		routing["rules"] = append(manualRules, managedRules...)
	}
	routing["balancers"] = append(manualBalancers, managedBalancers...)
	return updated, nil
}

func clone(value map[string]any) map[string]any {
	if value == nil {
		return map[string]any{}
	}
	b, _ := json.Marshal(value)
	var copy map[string]any
	_ = json.Unmarshal(b, &copy)
	return copy
}
func object(parent map[string]any, key string) map[string]any {
	if value, ok := parent[key].(map[string]any); ok {
		return value
	}
	value := map[string]any{}
	parent[key] = value
	return value
}
func objects(value any) ([]map[string]any, error) {
	if value == nil {
		return nil, nil
	}
	if list, ok := value.([]map[string]any); ok {
		return list, nil
	}
	list, ok := value.([]any)
	if !ok {
		return nil, errors.New("expected JSON array")
	}
	result := make([]map[string]any, 0, len(list))
	for _, item := range list {
		object, ok := item.(map[string]any)
		if !ok {
			return nil, errors.New("expected JSON object in array")
		}
		result = append(result, object)
	}
	return result, nil
}
func tags(value any) ([]string, error) {
	items, err := objects(value)
	if err != nil {
		return nil, err
	}
	result := []string{}
	for _, item := range items {
		if tag, ok := item["tag"].(string); ok {
			result = append(result, tag)
		}
	}
	return result, nil
}
func filterUnmanaged(items []map[string]any, key string) []map[string]any {
	result := []map[string]any{}
	for _, item := range items {
		value, _ := item[key].(string)
		if !strings.HasPrefix(value, managedPrefix) {
			result = append(result, item)
		}
	}
	return result
}
func matchingTags(tags []string, pattern string) []string {
	re := regexp.MustCompile(pattern)
	var result []string
	for _, tag := range tags {
		if re.MatchString(tag) {
			result = append(result, tag)
		}
	}
	sort.Strings(result)
	return result
}
func uniqueTags(tags []string) []string {
	set := map[string]bool{}
	for _, tag := range tags {
		if tag != "" {
			set[tag] = true
		}
	}
	result := make([]string, 0, len(set))
	for tag := range set {
		result = append(result, tag)
	}
	sort.Strings(result)
	return result
}
func validStrategy(value string) bool {
	return value == "random" || value == "roundRobin" || value == "leastPing" || value == "leastLoad"
}