package config

import "fmt"

type RuleAction string

const (
	RuleAccept   RuleAction = "accept"
	RuleDrop     RuleAction = "drop"
	RuleReject   RuleAction = "reject"
	RuleDNAT     RuleAction = "dnat"
	RuleRedirect RuleAction = "redirect"
	RuleLog      RuleAction = "log"
	RuleContinue RuleAction = "continue"
	RuleNFQueue  RuleAction = "nfqueue"
	RuleNoNAT    RuleAction = "nonat"
	RuleTarpit   RuleAction = "tarpit"
	RuleCount    RuleAction = "count"
	RuleMark     RuleAction = "mark"
	RuleConnMark RuleAction = "connmark"
)

type RuleSection string

const (
	SectionAll         RuleSection = "all"
	SectionEstablished RuleSection = "established"
	SectionRelated     RuleSection = "related"
	SectionInvalid     RuleSection = "invalid"
	SectionUntracked   RuleSection = "untracked"
	SectionNew         RuleSection = "new"
)

// Rule defines a specific traffic rule — an exception to the default policy.
// Rules are evaluated in order; the first terminating match wins.
// LOG, COUNT, MARK, and CONNMARK are non-terminating (packet continues to next rule).
type Rule struct {
	Action  RuleAction  `yaml:"action"`
	Section RuleSection `yaml:"section,omitempty"`

	// Source zone spec. Supports:
	//   zone, zone:address, zone:interface, zone:interface:address
	//   all, all+, any, none, all!zone1,zone2
	//   Multiple zones: loc,dmz
	Source string `yaml:"source"`

	// Dest zone spec. Same syntax as Source.
	// For DNAT: zone:server-ip:port[:random]
	// For REDIRECT: port (zone is implicitly the firewall)
	Dest string `yaml:"dest"`

	Proto string   `yaml:"proto,omitempty"`
	DPort PortSpec `yaml:"dport,omitempty"`
	SPort PortSpec `yaml:"sport,omitempty"`

	// PortGroup references a named portgroup (mutually exclusive with proto+dport).
	PortGroup string `yaml:"portgroup,omitempty"`

	Log string `yaml:"log,omitempty"`

	// OrigDest is the original destination address before DNAT/REDIRECT rewriting.
	// For non-NAT rules, constrains which original dest addresses match.
	OrigDest string `yaml:"origdest,omitempty"`

	// Rate limit. Format: [s:|d:][name:]rate/{sec|min|hour|day}[:burst]
	RateLimit string `yaml:"rate_limit,omitempty"`

	// User/group match (only valid when source is the firewall zone).
	// Format: [!]user[:group]
	User string `yaml:"user,omitempty"`

	// Packet or connection mark test. Format: [!]value[/mask][:C]
	Mark string `yaml:"mark,omitempty"`

	// Mark value to set (for mark/connmark actions). Format: value[/mask]
	SetMark string `yaml:"set_mark,omitempty"`

	// Simultaneous connection limit. Format: [d:]limit[:mask]
	ConnLimit string `yaml:"conn_limit,omitempty"`

	// Time-based restrictions.
	Time *TimeSpec `yaml:"time,omitempty"`

	// Conntrack helper. Values: ftp, sip, tftp, irc, pptp, amanda, snmp, etc.
	Helper string `yaml:"helper,omitempty"`

	// NFQUEUE number (only for nfqueue action).
	NFQueue int `yaml:"nfqueue,omitempty"`

	Comment string `yaml:"comment,omitempty"`
}

type TimeSpec struct {
	Start    string   `yaml:"start,omitempty"`
	Stop     string   `yaml:"stop,omitempty"`
	Weekdays []string `yaml:"weekdays,omitempty"`
	Monthdays []int   `yaml:"monthdays,omitempty"`
	DateStart string  `yaml:"date_start,omitempty"`
	DateStop  string  `yaml:"date_stop,omitempty"`
	UTC       bool    `yaml:"utc,omitempty"`
}

// PortSpec supports single ports, ranges, and lists.
// Examples: [80], [443], [80, 443], ["1024-65535"], [53, "80-90"]
// For ICMP, values are interpreted as ICMP types (e.g. "echo-request", "8", "3/4").
type PortSpec []string

func (ps *PortSpec) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var multi []interface{}
	if err := unmarshal(&multi); err == nil {
		for _, v := range multi {
			switch val := v.(type) {
			case int:
				*ps = append(*ps, fmt.Sprintf("%d", val))
			case float64:
				*ps = append(*ps, fmt.Sprintf("%d", int(val)))
			case string:
				*ps = append(*ps, val)
			default:
				return fmt.Errorf("unsupported port value type %T", v)
			}
		}
		return nil
	}

	var single string
	if err := unmarshal(&single); err == nil {
		*ps = PortSpec{single}
		return nil
	}

	var num int
	if err := unmarshal(&num); err == nil {
		*ps = PortSpec{fmt.Sprintf("%d", num)}
		return nil
	}

	return fmt.Errorf("invalid port spec")
}

var validRuleActions = map[RuleAction]bool{
	RuleAccept: true, RuleDrop: true, RuleReject: true,
	RuleDNAT: true, RuleRedirect: true, RuleLog: true,
	RuleContinue: true, RuleNFQueue: true, RuleNoNAT: true,
	RuleTarpit: true, RuleCount: true, RuleMark: true,
	RuleConnMark: true,
}

var validSections = map[RuleSection]bool{
	SectionAll: true, SectionEstablished: true, SectionRelated: true,
	SectionInvalid: true, SectionUntracked: true, SectionNew: true,
	"": true,
}

func (c *Config) validateRules() error {
	fwZone := c.FirewallZone()

	for i, r := range c.Rules {
		if !validRuleActions[r.Action] {
			return fmt.Errorf("rule[%d]: unknown action %q", i, r.Action)
		}
		if !validSections[r.Section] {
			return fmt.Errorf("rule[%d]: unknown section %q", i, r.Section)
		}

		if r.Source == "" {
			return fmt.Errorf("rule[%d]: source required", i)
		}
		if r.Dest == "" {
			return fmt.Errorf("rule[%d]: dest required", i)
		}

		if r.Source != "all" && r.Source != "any" && r.Source != "none" &&
			!hasPrefix(r.Source, "all+") && !hasPrefix(r.Source, "all!") && !hasPrefix(r.Source, "any!") {
			for _, srcPart := range splitZones(r.Source) {
				srcZone := zoneFromSpec(srcPart)
				if _, ok := c.Zones[srcZone]; !ok {
					return fmt.Errorf("rule[%d]: source zone %q not defined", i, srcZone)
				}
			}
		}

		if r.Action != RuleDNAT && r.Action != RuleRedirect && r.Action != RuleNoNAT {
			if r.Dest != "all" && r.Dest != "any" && r.Dest != "none" &&
				!hasPrefix(r.Dest, "all+") && !hasPrefix(r.Dest, "all!") && !hasPrefix(r.Dest, "any!") {
				for _, dstPart := range splitZones(r.Dest) {
					dstZone := zoneFromSpec(dstPart)
					if _, ok := c.Zones[dstZone]; !ok {
						return fmt.Errorf("rule[%d]: dest zone %q not defined", i, dstZone)
					}
				}
			}
		}

		if r.PortGroup != "" {
			if _, ok := c.PortGroups[r.PortGroup]; !ok {
				return fmt.Errorf("rule[%d]: portgroup %q not defined", i, r.PortGroup)
			}
			if r.Proto != "" || len(r.DPort) > 0 {
				return fmt.Errorf("rule[%d]: portgroup is mutually exclusive with proto/dport", i)
			}
		}

		if r.Action == RuleDNAT && r.Dest == "" {
			return fmt.Errorf("rule[%d]: dest with target address required for DNAT", i)
		}

		if r.User != "" && fwZone != "" {
			srcZone := zoneFromSpec(r.Source)
			if srcZone != fwZone {
				return fmt.Errorf("rule[%d]: user match only valid when source is the firewall zone", i)
			}
		}

		if (r.Action == RuleMark || r.Action == RuleConnMark) && r.SetMark == "" {
			return fmt.Errorf("rule[%d]: set_mark required for %s action", i, r.Action)
		}

		if r.Action == RuleTarpit && r.Proto != "tcp" && r.Proto != "" {
			return fmt.Errorf("rule[%d]: tarpit only works with proto tcp", i)
		}
	}
	return nil
}

// zoneFromSpec extracts the zone name from a zone spec like "net" or "net:192.168.1.0/24".
func zoneFromSpec(spec string) string {
	for i, c := range spec {
		if c == ':' {
			return spec[:i]
		}
	}
	return spec
}

func hasPrefix(s, prefix string) bool {
	return len(s) >= len(prefix) && s[:len(prefix)] == prefix
}