package ldap

import (
	"bytes"
	"errors"
	"fmt"

	ber "github.com/go-asn1-ber/asn1-ber"
)

const (
	maxFilterDepth     = 16
	maxFilterNodes     = 256
	maxAssertionBytes  = 1024
	maxAssertionsBytes = 16384
)

var (
	errMalformedFilter = errors.New("malformed filter")
	errFilterLimit     = errors.New("filter exceeds limits")
)

const (
	filterAnd            ber.Tag = 0
	filterOr             ber.Tag = 1
	filterNot            ber.Tag = 2
	filterEquality       ber.Tag = 3
	filterSubstrings     ber.Tag = 4
	filterGreaterOrEqual ber.Tag = 5
	filterLessOrEqual    ber.Tag = 6
	filterPresent        ber.Tag = 7
	filterApprox         ber.Tag = 8
	filterExtensible     ber.Tag = 9
)

const (
	substringInitial ber.Tag = 0
	substringAny     ber.Tag = 1
	substringFinal   ber.Tag = 2
)

const (
	extensibleRule  ber.Tag = 1
	extensibleType  ber.Tag = 2
	extensibleValue ber.Tag = 3
	extensibleDN    ber.Tag = 4
)

type FilterExpr struct {
	root *filterNode
}

type filterNode struct {
	kind         ber.Tag
	children     []*filterNode
	attr         string
	rule         *attrRule
	value        []byte
	substrings   []substring
	matchingRule string
	hasRule      bool
	dnAttributes bool
}

type substring struct {
	kind  ber.Tag
	value []byte
}

func Copy(p *ber.Packet) (FilterExpr, error) {
	var c copier
	root, err := c.filter(p, 1)
	if err != nil {
		return FilterExpr{}, err
	}
	return FilterExpr{root: root}, nil
}

type copier struct {
	nodes int
	bytes int
}

func malformed(format string, args ...any) error {
	return fmt.Errorf("%w: %s", errMalformedFilter, fmt.Sprintf(format, args...))
}

func limitExceeded(format string, args ...any) error {
	return fmt.Errorf("%w: %s", errFilterLimit, fmt.Sprintf(format, args...))
}

func (c *copier) filter(p *ber.Packet, depth int) (*filterNode, error) {
	if p == nil {
		return nil, malformed("missing filter")
	}
	if depth > maxFilterDepth {
		return nil, limitExceeded("depth exceeds %d", maxFilterDepth)
	}
	c.nodes++
	if c.nodes > maxFilterNodes {
		return nil, limitExceeded("more than %d nodes", maxFilterNodes)
	}
	if p.ClassType != ber.ClassContext {
		return nil, malformed("filter class %#x is not context-specific", uint8(p.ClassType))
	}
	switch p.Tag {
	case filterAnd, filterOr:
		return c.set(p, depth)
	case filterNot:
		if err := isConstructed(p, "not"); err != nil {
			return nil, err
		}
		if len(p.Children) != 1 {
			return nil, malformed("not has %d children", len(p.Children))
		}
		child, err := c.filter(p.Children[0], depth+1)
		if err != nil {
			return nil, err
		}
		return &filterNode{kind: filterNot, children: []*filterNode{child}}, nil
	case filterEquality, filterGreaterOrEqual, filterLessOrEqual, filterApprox:
		return c.assertion(p)
	case filterSubstrings:
		return c.substrings(p)
	case filterPresent:
		content, err := primitiveContent(p, ber.ClassContext, filterPresent, "present")
		if err != nil {
			return nil, err
		}
		attr, err := attributeDescription(content)
		if err != nil {
			return nil, err
		}
		return &filterNode{kind: filterPresent, attr: attr, rule: lookupAttribute(attr)}, nil
	case filterExtensible:
		return c.extensible(p)
	}
	return nil, malformed("unknown filter tag %d", uint64(p.Tag))
}

func (c *copier) set(p *ber.Packet, depth int) (*filterNode, error) {
	if err := isConstructed(p, "and/or"); err != nil {
		return nil, err
	}
	if len(p.Children) == 0 {
		return nil, malformed("empty and/or")
	}
	n := &filterNode{kind: p.Tag}
	for _, child := range p.Children {
		k, err := c.filter(child, depth+1)
		if err != nil {
			return nil, err
		}
		n.children = append(n.children, k)
	}
	return n, nil
}

func (c *copier) assertion(p *ber.Packet) (*filterNode, error) {
	if err := isConstructed(p, "attribute value assertion"); err != nil {
		return nil, err
	}
	if len(p.Children) != 2 {
		return nil, malformed("attribute value assertion has %d children", len(p.Children))
	}
	attr, err := describedAttribute(p.Children[0])
	if err != nil {
		return nil, err
	}
	content, err := primitiveContent(p.Children[1], ber.ClassUniversal, ber.TagOctetString, "assertion value")
	if err != nil {
		return nil, err
	}
	value, err := c.value(content)
	if err != nil {
		return nil, err
	}
	return &filterNode{kind: p.Tag, attr: attr, rule: lookupAttribute(attr), value: value}, nil
}

func (c *copier) substrings(p *ber.Packet) (*filterNode, error) {
	if err := isConstructed(p, "substrings"); err != nil {
		return nil, err
	}
	if len(p.Children) != 2 {
		return nil, malformed("substrings has %d children", len(p.Children))
	}
	attr, err := describedAttribute(p.Children[0])
	if err != nil {
		return nil, err
	}
	seq := p.Children[1]
	if seq == nil || seq.ClassType != ber.ClassUniversal || seq.TagType != ber.TypeConstructed || seq.Tag != ber.TagSequence {
		return nil, malformed("substrings components are not a SEQUENCE")
	}
	if len(seq.Children) == 0 {
		return nil, malformed("substrings without components")
	}
	n := &filterNode{kind: filterSubstrings, attr: attr, rule: lookupAttribute(attr)}
	last := len(seq.Children) - 1
	for i, comp := range seq.Children {
		if comp == nil || comp.ClassType != ber.ClassContext || comp.TagType != ber.TypePrimitive {
			return nil, malformed("substrings component %d is not a context-specific primitive", i)
		}
		switch {
		case comp.Tag == substringInitial && i != 0:
			return nil, malformed("substrings initial at position %d", i)
		case comp.Tag == substringFinal && i != last:
			return nil, malformed("substrings final at position %d", i)
		case comp.Tag > substringFinal:
			return nil, malformed("unknown substrings component tag %d", uint64(comp.Tag))
		}
		content, err := primitiveContent(comp, ber.ClassContext, comp.Tag, "substrings component")
		if err != nil {
			return nil, err
		}
		if len(content) == 0 {
			return nil, malformed("empty substrings component at position %d", i)
		}
		value, err := c.value(content)
		if err != nil {
			return nil, err
		}
		n.substrings = append(n.substrings, substring{kind: comp.Tag, value: value})
	}
	return n, nil
}

func (c *copier) extensible(p *ber.Packet) (*filterNode, error) {
	if err := isConstructed(p, "extensibleMatch"); err != nil {
		return nil, err
	}
	n := &filterNode{kind: filterExtensible}
	hasValue := false
	var previous ber.Tag
	for i, field := range p.Children {
		if field == nil {
			return nil, malformed("extensibleMatch field %d missing", i)
		}
		if field.Tag < extensibleRule || field.Tag > extensibleDN || (i > 0 && field.Tag <= previous) {
			return nil, malformed("extensibleMatch field tag %d out of order", uint64(field.Tag))
		}
		previous = field.Tag
		content, err := primitiveContent(field, ber.ClassContext, field.Tag, "extensibleMatch field")
		if err != nil {
			return nil, err
		}
		switch field.Tag {
		case extensibleRule:
			n.matchingRule = string(content)
			n.hasRule = true
		case extensibleType:
			attr, err := attributeDescription(content)
			if err != nil {
				return nil, err
			}
			n.attr = attr
		case extensibleValue:
			value, err := c.value(content)
			if err != nil {
				return nil, err
			}
			n.value = value
			hasValue = true
		case extensibleDN:
			if len(content) != 1 {
				return nil, malformed("extensibleMatch dnAttributes is %d bytes", len(content))
			}
			n.dnAttributes = content[0] != 0
		}
	}
	if !hasValue {
		return nil, malformed("extensibleMatch without matchValue")
	}
	if !n.hasRule && n.attr == "" {
		return nil, malformed("extensibleMatch without matchingRule or type")
	}
	return n, nil
}

func (c *copier) value(content []byte) ([]byte, error) {
	if len(content) > maxAssertionBytes {
		return nil, limitExceeded("assertion value of %d bytes exceeds %d", len(content), maxAssertionBytes)
	}
	c.bytes += len(content)
	if c.bytes > maxAssertionsBytes {
		return nil, limitExceeded("assertion values exceed %d bytes", maxAssertionsBytes)
	}
	return bytes.Clone(content), nil
}

func isConstructed(p *ber.Packet, what string) error {
	if p.TagType != ber.TypeConstructed {
		return malformed("%s is not constructed", what)
	}
	return nil
}

func primitiveContent(p *ber.Packet, class ber.Class, tag ber.Tag, what string) ([]byte, error) {
	if p == nil || p.ClassType != class || p.TagType != ber.TypePrimitive || p.Tag != tag {
		return nil, malformed("%s has the wrong tag", what)
	}
	if len(p.Children) != 0 || p.Data == nil {
		return nil, malformed("%s has no primitive content", what)
	}
	return p.Data.Bytes(), nil
}

func describedAttribute(p *ber.Packet) (string, error) {
	content, err := primitiveContent(p, ber.ClassUniversal, ber.TagOctetString, "attribute description")
	if err != nil {
		return "", err
	}
	return attributeDescription(content)
}

func attributeDescription(b []byte) (string, error) {
	parts := bytes.Split(b, []byte{';'})
	if !validOID(parts[0]) {
		return "", malformed("invalid attribute description")
	}
	for _, option := range parts[1:] {
		if len(option) == 0 || !allKeychars(option) {
			return "", malformed("invalid attribute description")
		}
	}
	return string(b), nil
}

func validOID(b []byte) bool {
	if len(b) == 0 {
		return false
	}
	if isAlpha(b[0]) {
		return allKeychars(b)
	}
	numbers := bytes.Split(b, []byte{'.'})
	if len(numbers) < 2 {
		return false
	}
	for _, number := range numbers {
		if len(number) == 0 || (len(number) > 1 && number[0] == '0') {
			return false
		}
		for _, d := range number {
			if !isDigit(d) {
				return false
			}
		}
	}
	return true
}

func allKeychars(b []byte) bool {
	for _, ch := range b {
		if !isAlpha(ch) && !isDigit(ch) && ch != '-' {
			return false
		}
	}
	return true
}

func isAlpha(ch byte) bool { return ('a' <= ch && ch <= 'z') || ('A' <= ch && ch <= 'Z') }

func isDigit(ch byte) bool { return '0' <= ch && ch <= '9' }

func (f FilterExpr) encode() *ber.Packet {
	if f.root == nil {
		return nil
	}
	return f.root.encode()
}

func (f FilterExpr) canonical() []byte {
	p := f.encode()
	if p == nil {
		return nil
	}
	return p.Bytes()
}

func (n *filterNode) encode() *ber.Packet {
	switch n.kind {
	case filterAnd, filterOr, filterNot:
		p := ber.Encode(ber.ClassContext, ber.TypeConstructed, n.kind, nil, "")
		for _, child := range n.children {
			p.AppendChild(child.encode())
		}
		return p
	case filterEquality, filterGreaterOrEqual, filterLessOrEqual, filterApprox:
		p := ber.Encode(ber.ClassContext, ber.TypeConstructed, n.kind, nil, "")
		p.AppendChild(octetString(n.attr))
		p.AppendChild(octetString(string(n.value)))
		return p
	case filterSubstrings:
		seq := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "")
		for _, s := range n.substrings {
			seq.AppendChild(ber.NewString(ber.ClassContext, ber.TypePrimitive, s.kind, string(s.value), ""))
		}
		p := ber.Encode(ber.ClassContext, ber.TypeConstructed, filterSubstrings, nil, "")
		p.AppendChild(octetString(n.attr))
		p.AppendChild(seq)
		return p
	case filterPresent:
		return ber.NewString(ber.ClassContext, ber.TypePrimitive, filterPresent, n.attr, "")
	}
	p := ber.Encode(ber.ClassContext, ber.TypeConstructed, filterExtensible, nil, "")
	if n.hasRule {
		p.AppendChild(ber.NewString(ber.ClassContext, ber.TypePrimitive, extensibleRule, n.matchingRule, ""))
	}
	if n.attr != "" {
		p.AppendChild(ber.NewString(ber.ClassContext, ber.TypePrimitive, extensibleType, n.attr, ""))
	}
	p.AppendChild(ber.NewString(ber.ClassContext, ber.TypePrimitive, extensibleValue, string(n.value), ""))
	if n.dnAttributes {
		p.AppendChild(ber.NewLDAPBoolean(ber.ClassContext, ber.TypePrimitive, extensibleDN, true, ""))
	}
	return p
}

func octetString(s string) *ber.Packet {
	return ber.NewString(ber.ClassUniversal, ber.TypePrimitive, ber.TagOctetString, s, "")
}