// Mnemonicode encoder/decoder.
// Encodes arbitrary bytes as human-readable words.
// Each group of 4 bytes maps to 3 words; trailing groups use fewer words.

package main

import (
  "embed"
  "encoding/json"
  "flag"
  "fmt"
  "io"
  "os"
  "path/filepath"
  "strings"
)

//go:embed wordlist.json
var wordlistJSON []byte

const mnBase = 1626

var (
  words     []string
  wordIndex map[string]int
)

func init() {
  var data struct {
    Words []string `json:"words"`
  }
  if err := json.Unmarshal(wordlistJSON, &data); err != nil {
    panic("failed to parse embedded wordlist: " + err.Error())
  }
  words = data.Words
  wordIndex = make(map[string]int, len(words))
  for i, w := range words {
    wordIndex[w] = i
  }
}

func wordsRequired(size int) int {
  return ((size + 1) * 3) / 4
}

func encode(data []byte) []string {
  result := make([]string, 0, wordsRequired(len(data)))
  for n := 0; n < wordsRequired(len(data)); n++ {
    grp := (n / 3) * 4
    remaining := len(data) - grp
    if remaining > 4 {
      remaining = 4
    }
    chunk := make([]byte, 4)
    copy(chunk, data[grp:grp+remaining])
    x := uint32(chunk[0]) | uint32(chunk[1])<<8 | uint32(chunk[2])<<16 | uint32(chunk[3])<<24
    pos := n % 3
    extra := 0
    if pos == 2 {
      if remaining == 3 {
        extra = mnBase
      }
      x /= mnBase * mnBase
    } else if pos == 1 {
      x /= mnBase
    }
    result = append(result, words[x%mnBase+uint32(extra)])
  }
  return result
}

func decode(encoded []string) ([]byte, error) {
  groups := len(encoded)/3 + 1
  buf := make([]byte, groups*4)
  offset := 0

  getU32 := func(pos int) uint32 {
    grp := pos & ^3
    return uint32(buf[grp]) | uint32(buf[grp+1])<<8 | uint32(buf[grp+2])<<16 | uint32(buf[grp+3])<<24
  }
  setU32 := func(pos int, v uint32) {
    grp := pos & ^3
    buf[grp] = byte(v)
    buf[grp+1] = byte(v >> 8)
    buf[grp+2] = byte(v >> 16)
    buf[grp+3] = byte(v >> 24)
  }

  for _, word := range encoded {
    idx, ok := wordIndex[strings.ToLower(word)]
    if !ok {
      return nil, fmt.Errorf("unrecognized word: %q", word)
    }

    pos := offset % 4
    x := getU32(offset)

    switch pos {
    case 3:
      return nil, fmt.Errorf("unexpected word after 24-bit remainder")
    case 0:
      x = uint32(idx)
      offset++
    case 1:
      x += uint32(idx) * mnBase
      offset++
    case 2:
      if idx >= mnBase {
        x += uint32(idx-mnBase) * mnBase * mnBase
        offset++ // offset%4 == 3 signals end-of-group
      } else {
        if idx > 1624 || (idx == 1624 && x > 1_312_671) {
          return nil, fmt.Errorf("invalid encoding: value out of range")
        }
        x += uint32(idx) * mnBase * mnBase
        offset += 2
      }
    }
    setU32(offset-1, x)
  }

  return buf[:offset], nil
}

func main() {
  sep := flag.String("s", "-", "word separator")
  flag.StringVar(sep, "separator", "-", "word separator")
  flag.Parse()

  decoding := strings.Contains(strings.ToLower(filepath.Base(os.Args[0])), "decode")

  if decoding {
    raw, err := io.ReadAll(os.Stdin)
    if err != nil {
      fmt.Fprintf(os.Stderr, "error reading stdin: %v\n", err)
      os.Exit(1)
    }
    parts := strings.FieldsFunc(string(raw), func(r rune) bool {
      return strings.ContainsRune(*sep, r) || r == '\n' || r == '\r'
    })
    var filtered []string
    for _, p := range parts {
      p = strings.TrimSpace(p)
      if p != "" {
        filtered = append(filtered, p)
      }
    }
    out, err := decode(filtered)
    if err != nil {
      fmt.Fprintf(os.Stderr, "error: %v\n", err)
      os.Exit(1)
    }
    os.Stdout.Write(out)
  } else {
    data, err := io.ReadAll(os.Stdin)
    if err != nil {
      fmt.Fprintf(os.Stderr, "error reading stdin: %v\n", err)
      os.Exit(1)
    }
    fmt.Println(strings.Join(encode(data), *sep))
  }
}