// Copyright 2024 Matthew Rich <matthewrich.conf@gmail.com>. All rights reserved.

package resource

import (
	"context"
_	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"gopkg.in/yaml.v3"
	"io"
	"net/url"
_	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"log/slog"
	"gitea.rosskeen.house/rosskeen.house/machine"
)

func init() {
	ResourceTypes.Register("iptable", func(u *url.URL) Resource {
		i := NewIptable()
		i.Table = IptableName(u.Hostname())
		if len(u.Path) > 0 {
			fields := strings.Split(u.Path, "/")
		slog.Info("iptables factory", "iptable", i, "uri", u, "fields", fields, "number_fields", len(fields))
			i.Chain = IptableChain(fields[1])
			if len(fields) < 3 {
				i.ResourceType = IptableTypeChain
			} else {
				i.ResourceType = IptableTypeRule
				id, _ := strconv.ParseUint(fields[2], 10, 32)
				i.Id = uint(id)
			}
			i.CreateCommand, i.ReadCommand, i.UpdateCommand, i.DeleteCommand = i.ResourceType.NewCRUD()
		}
		return i
	})
}

type IptableIPVersion string

const (
	IptableIPv4 IptableIPVersion = "ipv4"
	IPtableIPv6 IptableIPVersion = "ipv6"
)

type IptableName string

const (
	IptableNameFilter     = "filter"
	IptableNameNat        = "nat"
	IptableNameMangel     = "mangle"
	IptableNameRaw        = "raw"
	IptableNameSecurity   = "security"
)

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

type IptableChain string

const (
	IptableChainInput				= "INPUT"
	IptableChainOutput			= "OUTPUT"
	IptableChainForward			= "FORWARD"
	IptableChainPreRouting	= "PREROUTING"
	IptableChainPostRouting	= "POSTROUTING"
)

type IptableProto string

const (
	IptableProtoTCP = "tcp"
	IptableProtoUDP = "udp"
	IptableProtoUDPLite = "udplite"
	IptableProtoICMP = "icmp"
	IptableProtoICMPv6 = "icmpv6"
	IptableProtoESP = "ESP"
	IptableProtoAH = "AH"
	IptableProtoSCTP = "sctp"
	IptableProtoMH = "mh"
	IptableProtoAll = "all"
)

type IptableCIDR string

type ExtensionFlag struct {
	Name string `json:"name" yaml:"name"`
	Value string `json:"value" yaml:"value"`
}

type IptablePort uint16

type IptableType string

const (
	IptableTypeRule = "rule"
	IptableTypeChain = "chain"
)

// Manage the state of iptables rules
// iptable://filter/INPUT/0
type Iptable struct {
	Id						uint						`json:"id,omitempty" yaml:"id,omitempty"`
	Table					IptableName			`json:"table" yaml:"table"`
	Chain					IptableChain		`json:"chain" yaml:"chain"`
	Destination		IptableCIDR			`json:"destination,omitempty" yaml:"destination,omitempty"`
	Source				IptableCIDR			`json:"source,omitempty" yaml:"source,omitempty"`
	Dport					IptablePort			`json:"dport,omitempty" yaml:"dport,omitempty"`
	Sport					IptablePort			`json:"sport,omitempty" yaml:"sport,omitempty"`
	In						string					`json:"in,omitempty" yaml:"in,omitempty"`
	Out						string					`json:"out,omitempty" yaml:"out,omitempty"`
	Match					[]string				`json:"match,omitempty" yaml:"match,omitempty"`
	Flags					[]ExtensionFlag	`json:"extension_flags,omitempty" yaml:"extension_flags,omitempty"`
	Proto					IptableProto		`json:"proto,omitempty" yaml:"proto,omitempty"`
	Jump					string					`json:"jump,omitempty" yaml:"jump,omitempty"`
	State					string          `json:"state" yaml:"state"`
	ChainLength		uint						`json:"-" yaml:"-"`

	ResourceType	IptableType `json:"resourcetype,omitempty" yaml:"resourcetype,omitempty"`
  CreateCommand *Command `yaml:"-" json:"-"`
  ReadCommand   *Command `yaml:"-" json:"-"`
  UpdateCommand *Command `yaml:"-" json:"-"`
  DeleteCommand *Command `yaml:"-" json:"-"`
}

func NewIptable() *Iptable {
	i := &Iptable{ ResourceType: IptableTypeRule }
	i.CreateCommand, i.ReadCommand, i.UpdateCommand, i.DeleteCommand = i.ResourceType.NewCRUD()
	return i
}

func (i *Iptable) Clone() Resource {
	newIpt :=  &Iptable {
		Id: i.Id,
		Table: i.Table,
		Chain: i.Chain,
		Destination: i.Destination,
		Source: i.Source,
		In: i.In,
		Out: i.Out,
		Match: i.Match,
		Proto: i.Proto,
		ResourceType: i.ResourceType,
		State: i.State,
	}
	newIpt.CreateCommand, newIpt.ReadCommand, newIpt.UpdateCommand, newIpt.DeleteCommand = newIpt.ResourceType.NewCRUD()
	return newIpt
}

func (i *Iptable) StateMachine() machine.Stater {
	return StorageMachine()
}

func (i *Iptable) URI() string {
	return fmt.Sprintf("iptable://%s/%s/%d", i.Table, i.Chain, i.Id)
}

func (i *Iptable) SetURI(uri string) error {
	resourceUri, e := url.Parse(uri)
	if e == nil {
		if resourceUri.Scheme == "iptable" {
			i.Table = IptableName(resourceUri.Hostname())
			fields := strings.Split(resourceUri.Path, "/")
			i.Chain = IptableChain(fields[1])
			if len(fields) < 3 {
				i.ResourceType = IptableTypeChain
			} else {
				i.ResourceType = IptableTypeRule
				id, _ := strconv.ParseUint(fields[2], 10, 32)
				i.Id = uint(id)
			}
		} else {
			e = fmt.Errorf("%w: %s is not an iptable rule", ErrInvalidResourceURI, uri)
		}
	}
	return e
}

func (i *Iptable) Validate() error {
	s := NewSchema(i.Type())
	jsonDoc, jsonErr := i.JSON()
	if jsonErr == nil {
		return s.Validate(string(jsonDoc))
	}
	return jsonErr
}

func (i *Iptable) JSON() ([]byte, error) {
	return json.Marshal(i)
}

func (i *Iptable) UnmarshalJSON(data []byte) error {
	if unmarshalErr := json.Unmarshal(data, i); unmarshalErr != nil {
		return unmarshalErr
	}
	i.CreateCommand, i.ReadCommand, i.UpdateCommand, i.DeleteCommand = i.ResourceType.NewCRUD()
	return nil
}

func (i *Iptable) UnmarshalYAML(value *yaml.Node) error {
	type decodeIptable Iptable
	if unmarshalErr := value.Decode((*decodeIptable)(i)); unmarshalErr != nil {
		return unmarshalErr
	}
	i.CreateCommand, i.ReadCommand, i.UpdateCommand, i.DeleteCommand = i.ResourceType.NewCRUD()
	return nil
}

func (i *Iptable) NewCRUD() (create *Command, read *Command, update *Command, del *Command) {
	return NewIptableCreateCommand(), NewIptableReadCommand(), NewIptableUpdateCommand(), NewIptableDeleteCommand()
}

func (i *Iptable) Apply() error {
	ctx := context.Background()
	switch i.State {
	case "absent":
	case "present":
		err := i.Create(ctx)
		if err != nil {
			return err
		}
	}
	_,e := i.Read(context.Background())
	return e
}

func (i *Iptable) Load(r io.Reader) error {
	return NewYAMLDecoder(r).Decode(i)
}

func (i *Iptable) LoadDecl(yamlResourceDeclaration string) error {
	return NewYAMLStringDecoder(yamlResourceDeclaration).Decode(i)
}


func (i *Iptable) ResolveId(ctx context.Context) string {
//	uri := fmt.Sprintf("%s?gateway=%s&interface=%s&rtid=%s&metric=%d&type=%s&scope=%s",
//		n.To, n.Gateway, n.Interface, n.Rtid, n.Metric, n.RouteType, n.Scope)
//	n.Id = hex.EncodeToString([]byte(uri))
	return fmt.Sprintf("%d", i.Id)
}

func (i *Iptable) SetFlagValue(opt, value string) bool {
	switch opt {
	case "-i":
		i.In = value
		return true
	case "-o":
		i.Out = value
		return true
	case "-m":
		for _,search := range i.Match {
			if search == value {
				return true
			}
		}
		i.Match = append(i.Match, value)
		return true
	case "-s":
		i.Source = IptableCIDR(value)
		return true
	case "-d":
		i.Destination = IptableCIDR(value)
		return true
	case "-p":
		i.Proto = IptableProto(value)
		return true
	case "-j":
		i.Jump = value
		return true
	case "--dport":
		port,_ := strconv.ParseUint(value, 10, 16)
		i.Dport = IptablePort(port)
		return true
	case "--sport":
		port,_ := strconv.ParseUint(value, 10, 16)
		i.Sport = IptablePort(port)
		return true
	default:
    if opt[0] == '-' {
      i.Flags = append(i.Flags, ExtensionFlag{ Name: strings.Trim(opt, "-"), Value: strings.TrimSpace(value)})
			return true
    }
	}
	return false
}

func (i *Iptable) GetFlagValue(opt string) any {
	switch opt {
	case "-i":
		return i.In
	case "-o":
		return i.Out
	case "-m":
		return i.Match
	case "-s":
		return i.Source
	case "-d":
		return i.Destination
	case "-p":
		return i.Proto
	case "-j":
		return i.Jump
	case "--dport":
		return strconv.Itoa(int(i.Dport))
	case "--sport":
		return strconv.Itoa(int(i.Sport))
	default:
    if opt[0] == '-' {
      return i.Flags
    }
	}
	return nil
}

func (i *Iptable) SetRule(flags []string) (assigned bool) {
	assigned = true
	for index, flag := range flags {
		if flag[0] == '-' {
			flag := flags[index]
			value := flags[index + 1]
			if value[0] != '-' {
				if ! i.SetFlagValue(flag, value) {
					assigned = false
				}
			}
		}
	}
	return
}

func (i *Iptable) MatchRule(flags []string) (match bool) {
	match = true
	for index, flag := range flags {
		if flag[0] == '-' {
			value := flags[index + 1]
			switch v := i.GetFlagValue(flag).(type) {
			case []string:
				for _,element := range v {
					if element == value {
						continue
					}
				}
				match = false
			case []ExtensionFlag:
				for _,element := range v {
					if element.Name == flag && element.Value == value {
						continue
					}
				}
				match = false
			case IptableCIDR:
				if v == IptableCIDR(value) {
					continue
				}
				match = false
			case IptableName:
				if v == IptableName(value) {
					continue
				}
				match = false
			case IptableChain:
				if v == IptableChain(value) {
					continue
				}
				match = false
			default:
				if v.(string) == value {
					continue
				}
				match = false
			}
		}
	}
	return
}

func (i *Iptable) ReadChainLength() error {
	c := NewCommand()
	c.Path = "iptables"
	c.Args = []CommandArg{
		CommandArg("-S"),
		CommandArg("{{ .Chain }}"),
	}
	output,err := c.Execute(i)
	if err == nil {
		linesCount := strings.Count(string(output), "\n")
		if linesCount > 0 {
			i.ChainLength = uint(linesCount) - 1
		} else {
			i.ChainLength = 0
		}
	}
	return err
}

func (i *Iptable) Create(ctx context.Context) error {
	if i.Id > 0 {
		if lenErr := i.ReadChainLength(); lenErr != nil {
			return lenErr
		}
	}
	_, err := i.CreateCommand.Execute(i)
	//slog.Info("IptableChain Create()", "err", err, "errstr", err.Error(), "iptable", i, "createcommand", i.CreateCommand)
	// TODO add Command status/error handler rather than using the read extractor
	if i.CreateCommand.Extractor != nil {
		if err != nil {
			return i.CreateCommand.Extractor([]byte(err.Error()), i)
		}
	}
	return nil
}

func (i *Iptable) Read(ctx context.Context) ([]byte, error) {
	out, err := i.ReadCommand.Execute(i)
	if err != nil {
		return nil, err
	}
	exErr := i.ReadCommand.Extractor(out, i)
	if exErr != nil {
		return nil, exErr
	}
	return yaml.Marshal(i)
}

func (i *Iptable) Type() string { return "iptable" }

func (i *IptableType) NewCRUD() (create *Command, read *Command, update *Command, del *Command) {
	switch *i {
	case IptableTypeRule:
		return NewIptableCreateCommand(), NewIptableReadCommand(), NewIptableUpdateCommand(), NewIptableDeleteCommand()
	case IptableTypeChain:
		return NewIptableChainCreateCommand(), NewIptableChainReadCommand(), NewIptableChainUpdateCommand(), NewIptableChainDeleteCommand()
	default:
	}
	return nil, nil, nil, nil
}

func (i *IptableType) UnmarshalValue(value string) error {
	switch value {
	case string(IptableTypeRule), string(IptableTypeChain):
		*i = IptableType(value)
		return nil
	default:
		return errors.New("invalid IptableType value")
	}
}

func (i *IptableType) UnmarshalJSON(data []byte) error {
	var s string
	if unmarshalRouteTypeErr := json.Unmarshal(data, &s); unmarshalRouteTypeErr != nil {
		return unmarshalRouteTypeErr
	}
	return i.UnmarshalValue(s)
}

func (i *IptableType) UnmarshalYAML(value *yaml.Node) error {
	var s string
	if err := value.Decode(&s); err != nil {
		return err
	}
	return i.UnmarshalValue(s)
}

func NewIptableCreateCommand() *Command {
	c := NewCommand()
	c.Path = "iptables"
	c.Args = []CommandArg{
		CommandArg("-t"),
		CommandArg("{{ .Table }}"),
		CommandArg("{{ if le .Id .ChainLength }}-R{{ else }}-A{{ end }}"),
		CommandArg("{{ .Chain }}"),
		CommandArg("{{ if le .Id .ChainLength }}{{ .Id }}{{ end }}"),
		CommandArg("{{ if .In }}-i {{ .In }}{{ else if .Out }}-o {{ .Out }}{{ end }}"),
		CommandArg("{{ range .Match }}-m {{ . }} {{- end }}"),
		CommandArg("{{ if .Source }}-s {{ .Source }}{{ end }}"),
		CommandArg("{{ if .Sport }}--sport {{ .Sport }}{{ end }}"),
		CommandArg("{{ if .Destination }}-d {{ .Destination }}{{ end }}"),
		CommandArg("{{ if .Dport }}--dport {{ .Dport }}{{ end }}"),
		CommandArg("{{ if .Proto }}-p {{ .Proto }}{{ end }}"),
		CommandArg("{{ range .Flags }} --{{ .Name }} {{ .Value }} {{ end }}"),
		CommandArg("{{ if .Jump }}-j {{ .Jump }}{{ end }}"),
	}
	return c
}

func IptableExtractRule(lineNumber uint, ruleLine string, target *Iptable) (state string, err error) {
	state = "absent"
	ruleFields := strings.Split(strings.TrimSpace(ruleLine), " ")
	slog.Info("IptableExtractRule()", "lineNumber", lineNumber, "ruleLine", ruleLine, "target", target)
	if ruleFields[0] == "-A" {
		flags := ruleFields[2:]
		if target.Id > 0 {
			if target.Id == lineNumber {
				slog.Info("IptableExtractRule() SetRule", "lineNumber", lineNumber, "flags", flags, "target", target)
				if target.SetRule(flags) {
					state = "present"
					err = nil
				}
			}
		} else {
			if target.MatchRule(flags) {
				target.Id = lineNumber
				state = "present"
				err = nil
			}
		}
	} else {
		err = fmt.Errorf("Invalid rule %d %s", lineNumber, ruleLine)
	}
	return
}


func NewIptableReadCommand() *Command {
	c := NewCommand()
	c.Path = "iptables"
	c.Args = []CommandArg{
		CommandArg("-t"),
		CommandArg("{{ .Table }}"),
		CommandArg("-S"),
		CommandArg("{{ .Chain }}"),
		CommandArg("{{ if .Id }}{{ .Id }}{{ end }}"),
	}
	c.Extractor = func(out []byte, target any) error {
		i := target.(*Iptable)
		if i.Id > 0 {
			return RuleExtractor(out, target)
		}

		state := "absent"
		var lineNumber uint = 1
		lines := strings.Split(string(out), "\n")
		numberOfLines := len(lines)

		for _, line := range lines {
			matchState, err := IptableExtractRule(lineNumber, line, i)
			if matchState == "present" {
				state = matchState
				break
			}
			if err == nil {
				lineNumber++
			}
		}
		i.State = state
		if numberOfLines > 0 {
			i.ChainLength = uint(numberOfLines) - 1
		} else {
			i.ChainLength = 0
		}
		return nil
	}
	return c
}

func NewIptableReadChainCommand() *Command {
	c := NewCommand()
	c.Path = "iptables"
	c.Args = []CommandArg{
		CommandArg("-t"),
		CommandArg("{{ .Table }}"),
		CommandArg("-S"),
		CommandArg("{{ .Chain }}"),
	}
	c.Extractor = func(out []byte, target any) error {
		IptableChainRules := target.(*[]*Iptable)
		numberOfChainRules := len(*IptableChainRules)

		lines := strings.Split(strings.TrimSpace(string(out)), "\n")
		numberOfLines := len(lines)
		diff := (numberOfLines - 1) - numberOfChainRules
		if diff > 0 {
			for i := 0; i < diff; i++ {
				*IptableChainRules = append(*IptableChainRules, NewIptable())
			}
		}
		for lineIndex, line := range lines[1:] {
			i := (*IptableChainRules)[lineIndex]
			i.Id = uint(lineIndex + 1)
			ruleFields := strings.Split(strings.TrimSpace(line), " ")
			if ruleFields[0] == "-A" {
				flags := ruleFields[2:]
				if i.SetRule(flags) {
					i.State = "present"
				} else {
					i.State = "absent"
				}
			}
		}
		return nil
	}
	return c
}

func NewIptableUpdateCommand() *Command {
	return NewIptableCreateCommand()
}

func NewIptableDeleteCommand() *Command {
	return nil
}

func NewIptableChainCreateCommand() *Command {
	c := NewCommand()
	c.Path = "iptables"
	c.Args = []CommandArg{
		CommandArg("-t"),
		CommandArg("{{ .Table }}"),
		CommandArg("-N"),
		CommandArg("{{ .Chain }}"),
	}
	c.Extractor = func(out []byte, target any) error {
		slog.Info("IptableChain Extractor", "output", out, "command", c)
		for _,line := range strings.Split(string(out), "\n") {
			if line == "iptables: Chain already exists." {
				return nil
			}
		}
		return fmt.Errorf(string(out))
	}
	return c
}

func ChainExtractor(out []byte, target any) error {
	i := target.(*Iptable)
	rules := strings.Split(string(out), "\n")
	for _,rule := range rules {
		ruleFields := strings.Split(strings.TrimSpace(string(rule)), " ")
		switch ruleFields[0] {
		case "-N", "-A":
			chain := ruleFields[1]
			if chain == string(i.Chain) {
				i.State = "present"
				return nil
			} else {
				i.State = "absent"
			}
		default:
			i.State = "absent"
		}
	}
	return nil
}

func RuleExtractor(out []byte, target any) (err error) {
	ipt := target.(*Iptable)
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
	err = fmt.Errorf("Failed to extract rule by Id: %d", ipt.Id)
	ipt.State = "absent"
	var lineIndex uint = 1
	if uint(len(lines)) >= ipt.Id {
		lineIndex = ipt.Id
	} else if len(lines) > 2 {
		return
	}
	ruleFields := strings.Split(strings.TrimSpace(lines[lineIndex]), " ")
	slog.Info("RuleExtractor()", "lines", lines, "line", lines[lineIndex], "fields", ruleFields, "index", lineIndex)
	if ruleFields[0] == "-A" {
		if ipt.SetRule(ruleFields[2:]) {
			ipt.State = "present"
			err = nil
		}
	}
	return
}

func RuleExtractorMatchFlags(out []byte, target any) (err error) {
	ipt := target.(*Iptable)
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
	var linesCount uint = uint(len(lines))
	err = fmt.Errorf("Failed to extract rule")
	if linesCount > 0 {
		ipt.ChainLength = linesCount - 1
		ipt.State = "absent"
		for linesIndex, line := range lines {
			ruleFields := strings.Split(strings.TrimSpace(line), " ")
			slog.Info("RuleExtractorMatchFlags()", "lines", lines, "line", line, "fields", ruleFields, "index", linesIndex)
			if ruleFields[0] == "-A" {
				flags := ruleFields[2:]
				if ipt.MatchRule(flags) {
					slog.Info("RuleExtractorMatchFlags()", "flags", flags, "ipt", ipt)
					err = nil
					ipt.State = "present"
					ipt.Id = uint(linesIndex)
					return
				}
			}
		}
	}
	return
}

func RuleExtractorById(out []byte, target any) (err error) {
	ipt := target.(*Iptable)
	state := "absent"
	lines := strings.Split(string(out), "\n")
	err = fmt.Errorf("Failed to extract rule by Id: %d", ipt.Id)
	ipt.ChainLength = 0
	for _, line := range lines {
		ruleFields := strings.Split(strings.TrimSpace(line), " ")
		if ruleFields[0] == "-A" {
			ipt.ChainLength++
			flags := ruleFields[2:]
			slog.Info("RuleExtractorById()", "target", ipt)
			if ipt.Id == ipt.ChainLength {
				if ipt.SetRule(flags) {
					slog.Info("RuleExtractorById() SetRule", "flags", flags, "target", ipt)
					state = "present"
					err = nil
				}
			}
		}
	}
	ipt.State = state
	return
}

func NewIptableChainReadCommand() *Command {
	c := NewCommand()
	c.Path = "iptables"
	c.Args = []CommandArg{
		CommandArg("-t"),
		CommandArg("{{ .Table }}"),
		CommandArg("-S"),
		CommandArg("{{ .Chain }}"),
	}
	c.Extractor = func(out []byte, target any) error {
		i := target.(*Iptable)
		rules := strings.Split(string(out), "\n")
		for _,rule := range rules {
			ruleFields := strings.Split(strings.TrimSpace(string(rule)), " ")
			slog.Info("IptableChain Extract()", "fields", ruleFields)
			switch ruleFields[0] {
			case "-N", "-A":
				chain := ruleFields[1]
				if chain == string(i.Chain) {
					i.State = "present"
					return nil
				} else {
					i.State = "absent"
				}
			default:
				i.State = "absent"
			}
		}
		return nil
	}
	return c
}

func NewIptableChainUpdateCommand() *Command {
	return NewIptableChainCreateCommand()
}

func NewIptableChainDeleteCommand() *Command {
	return nil
}