package main

import (
	"errors"
	"fmt"
	"net/http"
	"net/netip"
	"regexp"
	"slices"
	"strconv"
	"strings"
)

// FirewallRule is an allow rule between two zones, as the page edits it:
// zones, addresses and ports on both sides, and the protocol. Rules with more
// (port forwards, rate limits, logging) are kept as they are and only listed.
type FirewallRule struct {
	From        string // zone
	Sources     string // addresses or CIDRs, empty: any
	SourcePorts string
	To          string // zone
	Destination string // one address or CIDR, empty: the whole zone
	Ports       string // destination ports
	Protocol    string
	Remove      bool
}

// ZoneMember is a LAN or a tunnel; its zone is what the firewall sees.
type ZoneMember struct {
	Kind   string // "lan" or "wireguard"
	Name   string
	Device string
	Zone   string
}

var portRangeRe = regexp.MustCompile(`^[0-9]+(-[0-9]+)?$`)

// simpleRuleKeys are the settings keys the page edits.
var simpleRuleKeys = []string{"from", "to", "protocol", "ports", "sourcePorts", "destination", "sources"}

// zoneMembers are the LANs, then the tunnels, by name. A member is in a zone
// named after itself unless it says otherwise.
func (s Settings) zoneMembers() []ZoneMember {
	router, _ := s["router"].(map[string]any)
	var out []ZoneMember
	lans := s.form().LANs
	for _, l := range lans {
		out = append(out, ZoneMember{Kind: "lan", Name: l.Name, Device: l.Interface, Zone: orName(str(router, "lans", l.Name, "zone"), l.Name)})
	}
	wgs, _ := router["wireguard"].(map[string]any)
	names := make([]string, 0, len(wgs))
	for n := range wgs {
		names = append(names, n)
	}
	slices.Sort(names)
	for _, n := range names {
		dev := orName(str(router, "wireguard", n, "interface"), "wg-"+n)
		out = append(out, ZoneMember{Kind: "wireguard", Name: n, Device: dev, Zone: orName(str(router, "wireguard", n, "zone"), n)})
	}
	return out
}

func orName(v, fallback string) string {
	if v == "" {
		return fallback
	}
	return v
}

// declaredZones are the zones listed in the settings (they may be empty).
func (s Settings) declaredZones() []string {
	router, _ := s["router"].(map[string]any)
	fw, _ := router["firewall"].(map[string]any)
	return stringList(fw, "zones")
}

// zones returns every firewall zone, sorted.
func (s Settings) zones() []string {
	out := slices.Clone(s.declaredZones())
	for _, m := range s.zoneMembers() {
		out = append(out, m.Zone)
	}
	slices.Sort(out)
	return slices.Compact(out)
}

func ruleSimple(r map[string]any) bool {
	if str(r, "from") == "wan" {
		return false
	}
	for k := range r {
		if !slices.Contains(simpleRuleKeys, k) {
			return false
		}
	}
	return true
}

func portsText(v any) string {
	var out []string
	items, _ := v.([]any)
	for _, i := range items {
		switch p := i.(type) {
		case string:
			out = append(out, p)
		case float64:
			out = append(out, strconv.FormatFloat(p, 'f', -1, 64))
		case int:
			out = append(out, strconv.Itoa(p))
		}
	}
	return strings.Join(out, ", ")
}

// firewallRules splits the rules into those the page edits and a description
// of the others, one line each.
func (s Settings) firewallRules() (simple []FirewallRule, other []string) {
	router, _ := s["router"].(map[string]any)
	fw, _ := router["firewall"].(map[string]any)
	rules, _ := fw["rules"].([]any)
	for _, v := range rules {
		r, _ := v.(map[string]any)
		if r == nil {
			continue
		}
		if wanRuleSimple(r) {
			continue
		}
		if !ruleSimple(r) {
			desc := str(r, "from") + " → " + str(r, "to")
			if d := str(r, "destination"); d != "" {
				desc += " " + d
			}
			if d := strings.Join(stringList(r, "sources"), ", "); d != "" {
				desc += " from " + d
			}
			if p := str(r, "protocol"); p != "" && p != "any" {
				desc += " " + p
			}
			if p := portsText(r["ports"]); p != "" {
				desc += " " + p
			}
			if dnat, _ := r["dnat"].(bool); dnat {
				desc += " (port forward)"
			}
			other = append(other, desc)
			continue
		}
		proto := str(r, "protocol")
		if proto == "" {
			proto = "any"
		}
		simple = append(simple, FirewallRule{From: str(r, "from"), To: str(r, "to"), Protocol: proto,
			Ports: portsText(r["ports"]), SourcePorts: portsText(r["sourcePorts"]),
			Sources: strings.Join(stringList(r, "sources"), ", "), Destination: str(r, "destination")})
	}
	return simple, other
}

// ZoneForm is what the page changes about the zones: the zone of each member,
// zones to add (separated by commas or spaces) and empty zones to drop.
type ZoneForm struct {
	Members []ZoneMember
	New     string
	Remove  []string
}

// portList parses "22, 443 5000-5100" into settings values (numbers, and
// strings for ranges).
func portList(v string) ([]any, string) {
	var ports []any
	for _, p := range splitList(v) {
		if !portRangeRe.MatchString(p) {
			return nil, "numbers or ranges like 5000-5100, separated by commas"
		}
		lo, hi, _ := strings.Cut(p, "-")
		l, _ := strconv.Atoi(lo)
		h := l
		if hi != "" {
			h, _ = strconv.Atoi(hi)
		}
		if l < 1 || h > 65535 || (hi != "" && h <= l) {
			return nil, p + " is not a valid port or range"
		}
		if hi == "" {
			ports = append(ports, l)
		} else {
			ports = append(ports, p)
		}
	}
	return ports, ""
}

func validAddress(v string) bool {
	if _, err := netip.ParsePrefix(v); err == nil {
		return true
	}
	_, err := netip.ParseAddr(v)
	return err == nil
}

// applyZones assigns the members' zones and declares or drops zones. Errors
// are keyed "member[i]" and "newzone".
func (s Settings) applyZones(zf ZoneForm, errs FormErrors) {
	router := obj(s, "router")
	used := map[string]bool{}
	for i, m := range zf.Members {
		m.Zone = strings.TrimSpace(m.Zone)
		switch {
		case !nameRe.MatchString(m.Zone):
			errs[fmt.Sprintf("member[%d]", i)] = "letters, digits and dashes, starting with a letter"
			continue
		case m.Zone == "wan":
			errs[fmt.Sprintf("member[%d]", i)] = "wan is the uplink's zone and cannot be assigned"
			continue
		}
		used[m.Zone] = true
		kind := map[string]string{"lan": "lans", "wireguard": "wireguard"}[m.Kind]
		members, _ := router[kind].(map[string]any)
		v, _ := members[m.Name].(map[string]any)
		if v == nil {
			continue // vanished since the page was drawn
		}
		if m.Zone == m.Name {
			delete(v, "zone")
		} else {
			v["zone"] = m.Zone
		}
	}
	declared := slices.Clone(s.declaredZones())
	for _, z := range splitList(zf.New) {
		switch {
		case !nameRe.MatchString(z) || z == "wan":
			errs["newzone"] = z + ": letters, digits and dashes, starting with a letter (wan is taken)"
		case !slices.Contains(declared, z):
			declared = append(declared, z)
		}
	}
	declared = slices.DeleteFunc(declared, func(z string) bool {
		if !slices.Contains(zf.Remove, z) {
			return false
		}
		if used[z] {
			errs["zone."+z] = "still has members"
			return false
		}
		return true
	})
	// What a member already names needs no declaration.
	declared = slices.DeleteFunc(declared, func(z string) bool { return used[z] })
	slices.Sort(declared)
	fw := obj(router, "firewall")
	if len(declared) == 0 {
		delete(fw, "zones")
	} else {
		list := make([]any, len(declared))
		for i, z := range declared {
			list[i] = z
		}
		fw["zones"] = list
	}
}

// applyFirewall changes the zones, then replaces the rules the page edits with
// rows (an untouched "add" row is skipped); the other rules stay as they are.
func (s Settings) applyFirewall(zf ZoneForm, rows []FirewallRule, fwds []WANRule) error {
	errs := FormErrors{}
	s.applyZones(zf, errs)
	zones := s.zones()
	var next []any
	seen := map[FirewallRule]bool{}
	for i, r := range rows {
		field := func(n string) string { return fmt.Sprintf("rule[%d].%s", i, n) }
		r.From, r.To = strings.TrimSpace(r.From), strings.TrimSpace(r.To)
		r.Ports, r.SourcePorts = strings.Join(splitList(r.Ports), ", "), strings.Join(splitList(r.SourcePorts), ", ")
		r.Sources, r.Destination = strings.Join(splitList(r.Sources), ", "), strings.TrimSpace(r.Destination)
		if r.Protocol == "" {
			r.Protocol = "any"
		}
		if r.Remove || (r.From == "" && r.To == "" && r.Ports == "" && r.SourcePorts == "" && r.Sources == "" && r.Destination == "") {
			continue
		}
		if !slices.Contains(zones, r.From) {
			errs[field("from")] = "choose a zone"
		}
		if !slices.Contains(zones, r.To) {
			errs[field("to")] = "choose a zone"
		}
		if r.From == r.To && r.From != "" {
			errs[field("to")] = "traffic inside a zone is already allowed"
		}
		if !slices.Contains([]string{"any", "tcp", "udp", "icmp"}, r.Protocol) {
			errs[field("protocol")] = "any, tcp, udp or icmp"
		}
		ports, msg := portList(r.Ports)
		if msg != "" {
			errs[field("ports")] = msg
		}
		sports, msg := portList(r.SourcePorts)
		if msg != "" {
			errs[field("sourcePorts")] = msg
		}
		if (len(ports) > 0 || len(sports) > 0) && r.Protocol != "tcp" && r.Protocol != "udp" {
			errs[field("ports")] = "ports need protocol tcp or udp"
		}
		var sources []any
		for _, a := range splitList(r.Sources) {
			if !validAddress(a) {
				errs[field("sources")] = a + " is not an address or CIDR"
				break
			}
			sources = append(sources, a)
		}
		if r.Destination != "" && !validAddress(r.Destination) {
			errs[field("destination")] = "one address or CIDR, or empty for the whole zone"
		}
		if seen[r] {
			errs[field("from")] = "the same rule twice"
		}
		seen[r] = true
		rule := map[string]any{"from": r.From, "to": r.To}
		if r.Protocol != "any" {
			rule["protocol"] = r.Protocol
		}
		if len(ports) > 0 {
			rule["ports"] = ports
		}
		if len(sports) > 0 {
			rule["sourcePorts"] = sports
		}
		if len(sources) > 0 {
			rule["sources"] = sources
		}
		if r.Destination != "" {
			rule["destination"] = r.Destination
		}
		next = append(next, rule)
	}
	next = append(next, s.forwardRules(fwds, zones, errs)...)
	router := obj(s, "router")
	fw := obj(router, "firewall")
	old, _ := fw["rules"].([]any)
	for _, v := range old {
		r, _ := v.(map[string]any)
		if r == nil || ruleSimple(r) || wanRuleSimple(r) {
			continue
		}
		// Port forwards and the like are not edited here, but must keep a zone to point at.
		if f := str(r, "from"); f != "wan" && !slices.Contains(zones, f) {
			errs["zone."+f] = "a rule that is not editable here still uses this zone"
		}
		if t := str(r, "to"); !slices.Contains(zones, t) {
			errs["zone."+t] = "a rule that is not editable here still uses this zone"
		}
		next = append(next, v)
	}
	if len(errs) > 0 {
		return errs
	}
	if len(next) == 0 {
		delete(fw, "rules")
	} else {
		fw["rules"] = next
	}
	if len(fw) == 0 {
		delete(router, "firewall")
	}
	return nil
}

type firewallData struct {
	Rows     []FirewallRule
	Forwards []WANRule
	Other    []string
	Zones    []string
	Members  []ZoneMember
	Empty    []string // zones without members, which can be removed
	New      string
	Errors   FormErrors
}

func (s Settings) firewallData(rows []FirewallRule, fwds []WANRule, members []ZoneMember, errs FormErrors) firewallData {
	_, other := s.firewallRules()
	d := firewallData{Rows: rows, Forwards: fwds, Other: other, Zones: s.zones(), Members: members, Errors: errs}
	if d.Members == nil {
		d.Members = s.zoneMembers()
	}
	for _, z := range s.declaredZones() {
		if !slices.ContainsFunc(s.zoneMembers(), func(m ZoneMember) bool { return m.Zone == z }) {
			d.Empty = append(d.Empty, z)
		}
	}
	return d
}

func (a *app) firewallPage(w http.ResponseWriter, r *http.Request, sess session) {
	p := page{Title: "Firewall", CSRF: sess.csrf, LoggedIn: true}
	s, err := loadSettings(a.sys.settingsPath())
	if err != nil {
		p.Error = err.Error()
		a.render(w, "firewall", p)
		return
	}
	rows, _ := s.firewallRules()
	add := FirewallRule{From: r.URL.Query().Get("from"), To: r.URL.Query().Get("to"), Protocol: "any"}
	p.Data = s.firewallData(append(rows, add), append(s.portForwards(), WANRule{Protocol: "tcp", NAT: true}), nil, nil)
	a.render(w, "firewall", p)
}

func (a *app) firewall(w http.ResponseWriter, r *http.Request, sess session) {
	p := page{Title: "Firewall", CSRF: sess.csrf, LoggedIn: true}
	s, err := loadSettings(a.sys.settingsPath())
	if err != nil {
		p.Error = err.Error()
		a.render(w, "firewall", p)
		return
	}
	zf := ZoneForm{New: r.PostFormValue("newzones"), Remove: r.PostForm["rmzone"]}
	nm, _ := strconv.Atoi(r.PostFormValue("members"))
	for i := range min(nm, 512) {
		v := func(k string) string { return r.PostFormValue(fmt.Sprintf("member.%d.%s", i, k)) }
		zf.Members = append(zf.Members, ZoneMember{Kind: v("kind"), Name: v("name"), Device: v("device"), Zone: v("zone")})
	}
	n, _ := strconv.Atoi(r.PostFormValue("rows"))
	nf, _ := strconv.Atoi(r.PostFormValue("forwards"))
	var fwds []WANRule
	for i := range min(nf, 512) {
		v := func(k string) string { return r.PostFormValue(fmt.Sprintf("fwd.%d.%s", i, k)) }
		fwds = append(fwds, WANRule{Protocol: v("protocol"), Ports: v("ports"), To: v("to"), Destination: v("destination"),
			ToPort: v("toport"), Sources: v("sources"), NAT: v("nat") == "on", Hairpin: v("hairpin") == "on",
			Log: v("log") == "on", RateLimit: v("ratelimit"), Remove: v("remove") == "on"})
	}
	var rows []FirewallRule
	for i := range min(n, 512) {
		v := func(k string) string { return r.PostFormValue(fmt.Sprintf("rule.%d.%s", i, k)) }
		rows = append(rows, FirewallRule{From: v("from"), Sources: v("sources"), SourcePorts: v("sourceports"),
			To: v("to"), Destination: v("destination"), Ports: v("ports"), Protocol: v("protocol"), Remove: v("remove") == "on"})
	}
	if err := s.applyFirewall(zf, rows, fwds); err != nil {
		var fe FormErrors
		errors.As(err, &fe)
		// Show what was entered; the zone lists are those before this attempt.
		s, _ = loadSettings(a.sys.settingsPath())
		p.Error = "Please correct the marked fields."
		d := s.firewallData(rows, fwds, zf.Members, fe)
		d.New = zf.New
		p.Data = d
		w.WriteHeader(http.StatusUnprocessableEntity)
		a.render(w, "firewall", p)
		return
	}
	if err := s.save(a.sys.settingsPath()); err != nil {
		p.Error = "Saving: " + err.Error()
		a.render(w, "firewall", p)
		return
	}
	back(w, r, "Saved. Apply to make the changes take effect.")
}