package kdbx

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"io"
	"strings"
)

// A KDBX file's payload is one XML document. It is kept here as a generic
// tree rather than a set of Go structs on purpose: a merge server that
// dropped elements it did not model would silently destroy data written by a
// newer KeePass. Every element, attribute and value survives a round trip,
// and the merge logic reaches into the few paths it actually understands.
type Node struct {
	Name  string
	Attrs []Attr
	Nodes []*Node
	// Text is the character data of a leaf element, already decrypted if the
	// element was protected.
	Text string
	// Protected records that this value was stored under the inner stream
	// cipher, so it is written back the same way.
	Protected bool
}

type Attr struct {
	Name  string
	Value string
}

// parseXML reads a KDBX payload into a tree.
func parseXML(data []byte) (*Node, error) {
	dec := xml.NewDecoder(bytes.NewReader(data))
	var stack []*Node
	var root *Node
	var text strings.Builder

	for {
		tok, err := dec.Token()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("kdbx: xml: %w", err)
		}
		switch t := tok.(type) {
		case xml.StartElement:
			n := &Node{Name: t.Name.Local}
			for _, a := range t.Attr {
				name := a.Name.Local
				if a.Name.Space != "" {
					name = a.Name.Space + ":" + a.Name.Local
				}
				if name == protectedAttr && strings.EqualFold(a.Value, "True") {
					n.Protected = true
					continue
				}
				n.Attrs = append(n.Attrs, Attr{Name: name, Value: a.Value})
			}
			if len(stack) > 0 {
				parent := stack[len(stack)-1]
				parent.Nodes = append(parent.Nodes, n)
			} else if root == nil {
				root = n
			}
			stack = append(stack, n)
			text.Reset()
		case xml.CharData:
			text.Write(t)
		case xml.EndElement:
			if len(stack) == 0 {
				return nil, fmt.Errorf("kdbx: xml: unbalanced </%s>", t.Name.Local)
			}
			n := stack[len(stack)-1]
			if len(n.Nodes) == 0 {
				n.Text = text.String()
			}
			stack = stack[:len(stack)-1]
			text.Reset()
		}
	}
	if root == nil {
		return nil, fmt.Errorf("kdbx: xml: empty document")
	}
	return root, nil
}

const protectedAttr = "Protected"

// marshalXML writes the tree back out, indented the way KeePass writes it.
func marshalXML(root *Node) []byte {
	var b bytes.Buffer
	b.WriteString(`<?xml version="1.0" encoding="utf-8" standalone="yes"?>` + "\n")
	writeNode(&b, root, 0)
	return b.Bytes()
}

func writeNode(b *bytes.Buffer, n *Node, depth int) {
	indent := strings.Repeat("\t", depth)
	b.WriteString(indent)
	b.WriteByte('<')
	b.WriteString(n.Name)
	for _, a := range n.Attrs {
		fmt.Fprintf(b, " %s=\"", a.Name)
		xml.EscapeText(b, []byte(a.Value))
		b.WriteByte('"')
	}
	if n.Protected {
		fmt.Fprintf(b, " %s=\"True\"", protectedAttr)
	}
	if len(n.Nodes) == 0 && n.Text == "" {
		b.WriteString("/>\n")
		return
	}
	b.WriteByte('>')
	if len(n.Nodes) == 0 {
		xml.EscapeText(b, []byte(n.Text))
		b.WriteString("</" + n.Name + ">\n")
		return
	}
	b.WriteByte('\n')
	for _, c := range n.Nodes {
		writeNode(b, c, depth+1)
	}
	b.WriteString(indent + "</" + n.Name + ">\n")
}

// ---- navigation helpers used by the merge logic ----

// Find returns the first child with this name, or nil.
func (n *Node) Find(name string) *Node {
	if n == nil {
		return nil
	}
	for _, c := range n.Nodes {
		if c.Name == name {
			return c
		}
	}
	return nil
}

// FindAll returns every child with this name.
func (n *Node) FindAll(name string) []*Node {
	if n == nil {
		return nil
	}
	out := []*Node{}
	for _, c := range n.Nodes {
		if c.Name == name {
			out = append(out, c)
		}
	}
	return out
}

// Path walks a chain of child names.
func (n *Node) Path(names ...string) *Node {
	cur := n
	for _, name := range names {
		cur = cur.Find(name)
		if cur == nil {
			return nil
		}
	}
	return cur
}

// Value is the text at a path, or "".
func (n *Node) Value(names ...string) string {
	if t := n.Path(names...); t != nil {
		return t.Text
	}
	return ""
}

// Set replaces, or creates, a direct child holding this text.
func (n *Node) Set(name, text string) *Node {
	if c := n.Find(name); c != nil {
		c.Text = text
		c.Nodes = nil
		return c
	}
	c := &Node{Name: name, Text: text}
	n.Nodes = append(n.Nodes, c)
	return c
}

// Remove drops a direct child node by identity.
func (n *Node) Remove(target *Node) bool {
	for i, c := range n.Nodes {
		if c == target {
			n.Nodes = append(n.Nodes[:i], n.Nodes[i+1:]...)
			return true
		}
	}
	return false
}

// Attr returns an attribute value.
func (n *Node) Attr(name string) string {
	if n == nil {
		return ""
	}
	for _, a := range n.Attrs {
		if a.Name == name {
			return a.Value
		}
	}
	return ""
}

// SetAttr sets or replaces an attribute.
func (n *Node) SetAttr(name, value string) {
	for i := range n.Attrs {
		if n.Attrs[i].Name == name {
			n.Attrs[i].Value = value
			return
		}
	}
	n.Attrs = append(n.Attrs, Attr{Name: name, Value: value})
}

// walk visits every node in document order, which is the order the inner
// stream cipher processes protected values in.
func (n *Node) walk(fn func(*Node)) {
	if n == nil {
		return
	}
	fn(n)
	for _, c := range n.Nodes {
		c.walk(fn)
	}
}

// clone deep-copies a subtree, so a node taken from an incoming database does
// not stay aliased to it.
func (n *Node) clone() *Node {
	if n == nil {
		return nil
	}
	out := &Node{Name: n.Name, Text: n.Text, Protected: n.Protected}
	if len(n.Attrs) > 0 {
		out.Attrs = append([]Attr(nil), n.Attrs...)
	}
	for _, c := range n.Nodes {
		out.Nodes = append(out.Nodes, c.clone())
	}
	return out
}