"""
Pure-Python implementation of mnemonicode encoding/decoding.

Mnemonicode encodes arbitrary bytes as a sequence of human-readable words.
Each group of 4 bytes maps to 3 words; shorter trailing groups use fewer words.
The 24-bit trailing case (3 bytes) uses a distinct set of 7 short words to
signal that a remainder word ends the group.

Word list: version 0.7, 1626 base words + 7 remainder words = 1633 total.
"""

import json
import sys
from functools import lru_cache
from pathlib import Path

MN_BASE = 1626

_WORDLIST_JSON = Path(__file__).parent / "wordlist.json"


@lru_cache(maxsize=1)
def _load() -> tuple[list[str], dict[str, int], str]:
  data = json.loads(_WORDLIST_JSON.read_text())
  words: list[str] = data["words"]
  index: dict[str, int] = {w: i for i, w in enumerate(words)}
  return words, index, data["version"]


def wordlist_version() -> str:
  return _load()[2]


def words_required(size: int) -> int:
  """Number of mnemonic words needed to encode `size` bytes."""
  return ((size + 1) * 3) // 4


def encode(data: bytes) -> list[str]:
  """Encode bytes to a list of mnemonic words."""
  words, _, _ = _load()
  result = []
  for n in range(words_required(len(data))):
    grp = (n // 3) * 4
    remaining = min(len(data) - grp, 4)
    x = int.from_bytes(data[grp:grp + remaining], "little")
    pos = n % 3
    extra = 0
    if pos == 2:
      if remaining == 3:
        extra = MN_BASE  # use one of the 7 short remainder words
      x //= MN_BASE * MN_BASE
    elif pos == 1:
      x //= MN_BASE
    result.append(words[x % MN_BASE + extra])
  return result


def decode(encoded: list[str]) -> bytes:
  """Decode a list of mnemonic words back to bytes."""
  _, word_index, _ = _load()

  buf = bytearray(((len(encoded) // 3) + 1) * 4)
  offset = 0

  for word in encoded:
    try:
      idx = word_index[word.lower()]
    except KeyError:
      raise ValueError(f"unrecognized word: {word!r}") from None

    grp = offset & ~3
    x = int.from_bytes(buf[grp:grp + 4], "little")
    pos = offset % 4

    if pos == 3:
      raise ValueError("unexpected word after 24-bit remainder")
    elif pos == 0:
      x = idx
      offset += 1
    elif pos == 1:
      x += idx * MN_BASE
      offset += 1
    elif pos == 2:
      if idx >= MN_BASE:  # 24-bit remainder word
        x += (idx - MN_BASE) * MN_BASE * MN_BASE
        offset += 1  # offset % 4 == 3 signals end-of-group next
      else:
        if idx > 1624 or (idx == 1624 and x > 1_312_671):
          raise ValueError("invalid encoding: value out of range")
        x += idx * MN_BASE * MN_BASE
        offset += 2  # skip to next group

    buf[grp:grp + 4] = x.to_bytes(4, "little")

  return bytes(buf[:offset])


if __name__ == "__main__":
  import argparse

  parser = argparse.ArgumentParser(
    prog="mnemonicode",
    description="Encode bytes to mnemonic words or decode words back to bytes.",
  )
  parser.add_argument("-s", "--separator", default="-", help="word separator (default: -)")
  sub = parser.add_subparsers(dest="cmd", required=True)
  sub.add_parser("encode", help="read raw bytes from stdin, write words to stdout")
  sub.add_parser("decode", help="read words from stdin, write raw bytes to stdout")

  args = parser.parse_args()

  if args.cmd == "encode":
    data = sys.stdin.buffer.read()
    print(args.separator.join(encode(data)))

  elif args.cmd == "decode":
    words = [w.strip() for w in sys.stdin.read().split(args.separator) if w.strip()]
    try:
      data = decode(words)
    except ValueError as e:
      print(f"error: {e}", file=sys.stderr)
      sys.exit(1)
    sys.stdout.buffer.write(data)