// Copyright 2015 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:generate stringer -output binres_string.go -type ResType,DataType

// Package binres implements encoding and decoding of android binary resources.
//
// Binary resource structs support unmarshalling the binary output of aapt.
// Implementations of marshalling for each struct must produce the exact input
// sent to unmarshalling. This allows tests to validate each struct representation
// of the binary format as follows:
//
//  * unmarshal the output of aapt
//  * marshal the struct representation
//  * perform byte-to-byte comparison with aapt output per chunk header and body
//
// This process should strive to make structs idiomatic to make parsing xml text
// into structs trivial.
//
// Once the struct representation is validated, tests for parsing xml text
// into structs can become self-referential as the following holds true:
//
//  * the unmarshalled input of aapt output is the only valid target
//  * the unmarshalled input of xml text may be compared to the unmarshalled
//    input of aapt output to identify errors, e.g. text-trims, wrong flags, etc
//
// This provides validation, byte-for-byte, for producing binary xml resources.
//
// It should be made clear that unmarshalling binary resources is currently only
// in scope for proving that the BinaryMarshaler works correctly. Any other use
// is currently out of scope.
//
// A simple view of binary xml document structure:
//
//  XML
//    Pool
//    Map
//    Namespace
//    [...node]
//
// Additional resources:
// https://android.googlesource.com/platform/frameworks/base/+/master/include/androidfw/ResourceTypes.h
// https://justanapplication.wordpress.com/2011/09/13/ (a series of articles, increment date)
package binres

import (
	"encoding"
	"encoding/binary"
	"fmt"
)

type ResType uint16

func (t ResType) IsSupported() bool {
	// explicit for clarity
	return t == ResStringPool || t == ResXML ||
		t == ResXMLStartNamespace || t == ResXMLEndNamespace ||
		t == ResXMLStartElement || t == ResXMLEndElement ||
		t == ResXMLCharData ||
		t == ResXMLResourceMap ||
		t == ResTable || t == ResTablePackage
}

// explicitly defined for clarity and resolvability with apt source
const (
	ResNull       ResType = 0x0000
	ResStringPool ResType = 0x0001
	ResTable      ResType = 0x0002
	ResXML        ResType = 0x0003

	ResXMLStartNamespace ResType = 0x0100
	ResXMLEndNamespace   ResType = 0x0101
	ResXMLStartElement   ResType = 0x0102
	ResXMLEndElement     ResType = 0x0103
	ResXMLCharData       ResType = 0x0104

	ResXMLResourceMap ResType = 0x0180

	ResTablePackage  ResType = 0x0200
	ResTableType     ResType = 0x0201
	ResTableTypeSpec ResType = 0x0202
)

var (
	btou16 = binary.LittleEndian.Uint16
	btou32 = binary.LittleEndian.Uint32
	putu16 = binary.LittleEndian.PutUint16
	putu32 = binary.LittleEndian.PutUint32
)

// unmarshaler wraps BinaryUnmarshaler to provide byte size of decoded chunks.
type unmarshaler interface {
	encoding.BinaryUnmarshaler

	// size returns the byte size unmarshalled after a call to
	// UnmarshalBinary, or otherwise zero.
	size() int
}

// chunkHeader appears at the front of every data chunk in a resource.
// TODO look into removing this, it's not necessary for marshalling and
// the information provided may possibly be given more simply. For unmarshal,
// a simple function would do.
type chunkHeader struct {
	// Type of data that follows this header.
	typ ResType

	// Advance slice index by this value to find its associated data, if any.
	headerByteSize uint16

	// This is the header size plus the size of any data associated with the chunk.
	// Advance slice index by this value to completely skip its contents, including
	// any child chunks. If this value is the same as headerByteSize, there is
	// no data associated with the chunk.
	byteSize uint32
}

// size implements unmarshaler.
func (hdr chunkHeader) size() int { return int(hdr.byteSize) }

func (hdr *chunkHeader) UnmarshalBinary(bin []byte) error {
	hdr.typ = ResType(btou16(bin))
	if !hdr.typ.IsSupported() {
		return fmt.Errorf("%s not supported", hdr.typ)
	}
	hdr.headerByteSize = btou16(bin[2:])
	hdr.byteSize = btou32(bin[4:])
	return nil
}

func (hdr chunkHeader) MarshalBinary() ([]byte, error) {
	if !hdr.typ.IsSupported() {
		return nil, fmt.Errorf("%s not supported", hdr.typ)
	}
	bin := make([]byte, 8)
	putu16(bin, uint16(hdr.typ))
	putu16(bin[2:], hdr.headerByteSize)
	putu32(bin[4:], hdr.byteSize)
	return bin, nil
}

type XML struct {
	chunkHeader

	Pool *Pool
	Map  *Map

	Namespace *Namespace
	Children  []*Element

	// tmp field used when unmarshalling
	stack []*Element
}

// TODO this is used strictly for querying in tests and dependent on
// current XML.UnmarshalBinary implementation. Look into moving directly
// into tests.
var debugIndices = make(map[encoding.BinaryMarshaler]int)

func (bx *XML) UnmarshalBinary(bin []byte) error {
	buf := bin
	if err := (&bx.chunkHeader).UnmarshalBinary(bin); err != nil {
		return err
	}
	buf = buf[8:]

	// TODO this is tracked strictly for querying in tests; look into moving this
	// functionality directly into tests if possible.
	debugIndex := 8

	for len(buf) > 0 {
		t := ResType(btou16(buf))
		k, err := bx.kind(t)
		if err != nil {
			return err
		}
		if err := k.UnmarshalBinary(buf); err != nil {
			return err
		}
		debugIndices[k.(encoding.BinaryMarshaler)] = debugIndex
		debugIndex += int(k.size())
		buf = buf[k.size():]
	}
	return nil
}

func (bx *XML) kind(t ResType) (unmarshaler, error) {
	switch t {
	case ResStringPool:
		if bx.Pool != nil {
			return nil, fmt.Errorf("pool already exists")
		}
		bx.Pool = new(Pool)
		return bx.Pool, nil
	case ResXMLResourceMap:
		if bx.Map != nil {
			return nil, fmt.Errorf("resource map already exists")
		}
		bx.Map = new(Map)
		return bx.Map, nil
	case ResXMLStartNamespace:
		if bx.Namespace != nil {
			return nil, fmt.Errorf("namespace start already exists")
		}
		bx.Namespace = new(Namespace)
		return bx.Namespace, nil
	case ResXMLEndNamespace:
		if bx.Namespace.end != nil {
			return nil, fmt.Errorf("namespace end already exists")
		}
		bx.Namespace.end = new(Namespace)
		return bx.Namespace.end, nil
	case ResXMLStartElement:
		el := new(Element)
		if len(bx.stack) == 0 {
			bx.Children = append(bx.Children, el)
		} else {
			n := len(bx.stack)
			var p *Element
			p, bx.stack = bx.stack[n-1], bx.stack[:n-1]
			p.Children = append(p.Children, el)
			bx.stack = append(bx.stack, p)
		}
		bx.stack = append(bx.stack, el)
		return el, nil
	case ResXMLEndElement:
		n := len(bx.stack)
		var el *Element
		el, bx.stack = bx.stack[n-1], bx.stack[:n-1]
		if el.end != nil {
			return nil, fmt.Errorf("element end already exists")
		}
		el.end = new(ElementEnd)
		return el.end, nil
	case ResXMLCharData: // TODO
		cdt := new(CharData)
		el := bx.stack[len(bx.stack)-1]
		if el.head == nil {
			el.head = cdt
		} else if el.tail == nil {
			el.tail = cdt
		} else {
			return nil, fmt.Errorf("element head and tail already contain chardata")
		}
		return cdt, nil
	default:
		return nil, fmt.Errorf("unexpected type %s", t)
	}
}

func (bx *XML) MarshalBinary() ([]byte, error) {
	var (
		bin, b []byte
		err    error
	)
	b, err = bx.chunkHeader.MarshalBinary()
	if err != nil {
		return nil, err
	}
	bin = append(bin, b...)

	b, err = bx.Pool.MarshalBinary()
	if err != nil {
		return nil, err
	}
	bin = append(bin, b...)

	b, err = bx.Map.MarshalBinary()
	if err != nil {
		return nil, err
	}
	bin = append(bin, b...)

	b, err = bx.Namespace.MarshalBinary()
	if err != nil {
		return nil, err
	}
	bin = append(bin, b...)

	for _, child := range bx.Children {
		if err := marshalRecurse(child, &bin); err != nil {
			return nil, err
		}
	}

	b, err = bx.Namespace.end.MarshalBinary()
	if err != nil {
		return nil, err
	}
	bin = append(bin, b...)

	return bin, nil
}

func marshalRecurse(el *Element, bin *[]byte) error {
	b, err := el.MarshalBinary()
	if err != nil {
		return err
	}
	*bin = append(*bin, b...)

	if el.head != nil {
		b, err := el.head.MarshalBinary()
		if err != nil {
			return err
		}
		*bin = append(*bin, b...)
	}

	for _, child := range el.Children {
		if err := marshalRecurse(child, bin); err != nil {
			return err
		}
	}

	b, err = el.end.MarshalBinary()
	if err != nil {
		return err
	}
	*bin = append(*bin, b...)

	return nil
}

func (bx *XML) iterElements() <-chan *Element {
	ch := make(chan *Element, 1)
	go func() {
		for _, el := range bx.Children {
			iterElementsRecurse(el, ch)
		}
		close(ch)
	}()
	return ch
}

func iterElementsRecurse(el *Element, ch chan *Element) {
	ch <- el
	for _, e := range el.Children {
		iterElementsRecurse(e, ch)
	}
}