package main

import (
    "errors"
    "fmt"
    "io/ioutil"
    "strings"
)

type planet struct {
    name     string
    parent   *planet
    children []*planet
}

func main() {
    inputFile, _ := ioutil.ReadFile("input.txt")
    orbits := strings.Split(string(inputFile), "\n")

    rootPlanet := planet{name: "COM"}
    rootPlanet.addOrbits(orbits)
    fmt.Println("Part 1:", rootPlanet.countOrbits(0))
    fmt.Println("Part 2:", rootPlanet.findDistance("YOU", "SAN"))
}

func (p *planet) addOrbits(relationships []string) []string {
    var i int
    for i < len(relationships) {
        ab := strings.Split(relationships[i], ")")
        if ab[0] == p.name {
            child := &planet{name: ab[1], parent: p}
            p.children = append(p.children, child)
            relationships = remove(relationships, i)
        } else {
            i++
        }
    }
    for _, child := range p.children {
        relationships = child.addOrbits(relationships)
    }
    return relationships
}

func remove(a []string, i int) []string {
    a[i] = a[len(a)-1]  // Copy last element to index i.
    a[len(a)-1] = ""    // Erase last element (write zero value).
    return a[:len(a)-1] // Truncate slice.
}

func (p *planet) countOrbits(count int) int {
    var additional int
    for _, child := range p.children {
        additional += child.countOrbits(count + 1)
    }
    return count + additional
}

func (p *planet) findObject(name string) (int, error) {
    if p.name == name {
        return -1, nil
    }
    for _, child := range p.children {
        distance, err := child.findObject(name)
        if err == nil {
            return distance + 1, nil
        }
    }

    return 0, errors.New("not found")
}

func (p *planet) findDistance(a, b string) int {
    d1, err1 := p.findObject(a)
    d2, err2 := p.findObject(b)
    if err1 != nil || err2 != nil {
        return -1
    }
    for _, child := range p.children {
        d := child.findDistance(a, b)
        if d > 0 {
            return d
        }
    }
    return d1 + d2
}