notea-de-me/spike/review/reinvent/ezmded/server/tree.go

package main

import (
	"errors"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"time"

	yaml "gopkg.in/yaml.v2"
)

var errChainNotFound = errors.New("pid chain not found")

type Tree struct {
	root string
}

type Branch struct {
	Title   string
	Deleted bool
	Updated time.Time
	PID     string
}

type Pretty struct {
	Children map[string]*Pretty
	Title    string
}

func NewTree(root string) *Tree {
	return &Tree{
		root: root,
	}
}

func (tree *Tree) GetPretty() (map[string]*Pretty, error) {
	branches, err := tree.Get()
	if err != nil {
		return nil, err
	}
	_ = branches
	realpretty := map[string]*Pretty{}
	lookuppretty := map[string]*Pretty{}
	timesZero := 0
	for len(branches) > 0 && timesZero < 3 {
		topop := []string{}
		for id, branch := range branches {
			if branch.Deleted {
			} else if branch.PID == "" {
				realpretty[id] = &Pretty{
					Title:    branch.Title,
					Children: map[string]*Pretty{},
				}
				lookuppretty[id] = realpretty[id]
			} else if ptr, ok := lookuppretty[branch.PID]; ok {
				ptr.Children[id] = &Pretty{
					Title:    branch.Title,
					Children: map[string]*Pretty{},
				}
				lookuppretty[id] = ptr.Children[id]
			} else {
				continue
			}
			topop = append(topop, id)
		}
		for _, id := range topop {
			delete(branches, id)
		}
		if len(topop) == 0 {
			timesZero++
		} else {
			timesZero = 0
		}
	}
	return realpretty, nil
}

func (tree *Tree) Get() (map[string]Branch, error) {
	m := map[string]Branch{}
	err := filepath.Walk(tree.root, func(p string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		if info.Name() != "meta.yaml" {
			return nil
		}
		b, err := ioutil.ReadFile(p)
		if err != nil {
			return err
		}
		var branch Branch
		if err := yaml.Unmarshal(b, &branch); err != nil {
			return err
		}
		id := path.Base(path.Dir(p))
		pidFullPath := path.Dir(path.Dir(p))
		if pidFullPath != tree.root {
			branch.PID = path.Base(pidFullPath)
		} else {
			branch.PID = ""
		}
		m[id] = branch
		return nil
	})
	if os.IsNotExist(err) {
		return m, nil
	}
	return m, err
}

func (tree *Tree) fullIdAndMeta(id string) (string, Branch, error) {
	m, err := tree.Get()
	if err != nil {
		return "", Branch{}, err
	}
	branch, ok := m[id]
	if !ok {
		return "", Branch{}, errChainNotFound
	}
	p := id
	for branch.PID != "" {
		p = path.Join(branch.PID, p)
		branch, ok = m[branch.PID]
		if !ok {
			return "", Branch{}, errChainNotFound
		}
	}
	return path.Join(tree.root, p), branch, nil
}

func (tree *Tree) putMeta(fullId string, branch Branch) error {
	b, err := yaml.Marshal(branch)
	if err != nil {
		return err
	}
	return ensureAndWrite(path.Join(fullId, "meta.yaml"), b)
}

func (tree *Tree) Del(id string) error {
	fullId, meta, err := tree.fullIdAndMeta(id)
	if err != nil {
		if err == errChainNotFound {
			return nil
		}
		return err
	}

	meta.Updated = time.Now().UTC()
	meta.Deleted = true

	return tree.putMeta(fullId, meta)
}

func (tree *Tree) Put(id string, branch Branch) error {
	fullId, meta, err := tree.fullIdAndMeta(id)
	if err == errChainNotFound {
		if branch.PID != "" {
			fullId, _, err = tree.fullIdAndMeta(branch.PID)
			fullId = path.Join(fullId, id)
		} else {
			err = nil
			fullId = path.Join(tree.root, fullId, id)
		}
	}
	if err != nil {
		return err
	}

	meta.Updated = time.Now().UTC()
	if meta.Title == "" {
		meta.Title = "Untitled (" + time.Now().UTC().String() + ")"
	}
	if branch.Title != "" {
		meta.Title = branch.Title
	}
	meta.PID = branch.PID
	meta.Deleted = branch.Deleted

	return tree.putMeta(fullId, meta)
}