simple-nanoshare/link_composer/codec.py

from dataclasses import dataclass
from urllib.parse import urlsplit, urlunsplit
import re
import zlib

from .constants import (
  ALPHABET_NAMES,
  CODE_PREFIX_PAYLOAD,
  CODE_PREFIX_QR,
  CODE_SCHEMES,
  CODE_VERSION,
  DOMAIN_CODE_WIDTH,
  LEGACY_DICTIONARY_CODE_WIDTH,
  COMMON_DOMAINS,
  COMMON_TLDS,
  MODE_QR_STRUCTURED,
  MODE_STRUCTURED,
  PAYLOAD_ALPHABET,
  QR_ALPHABET,
  SCHEME_CODES,
  SUPPORTED_CODE_VERSIONS,
  SEGMENT_ALPHABETS,
  SEPARATORS,
  TLD_CODE_WIDTH,
)
from .base_n import bytes_to_int, decode_int, encode_int, int_to_bytes
from .bitstream import BitReader, BitWriter

@dataclass(frozen=True)
class EncodeResult:
  code: str
  original_length: int
  code_length: int
  bit_length: int
  mode: str

def encode_url(url: str, *, qr: bool = False) -> EncodeResult:
  normalized = _normalize_url(url)
  writer = BitWriter()
  writer.write(CODE_VERSION, 4)

  writer.write(0, 1)
  _write_structured_url(writer, normalized)

  if qr:
    alphabet = QR_ALPHABET
    mode = MODE_QR_STRUCTURED
  else:
    alphabet = PAYLOAD_ALPHABET
    mode = MODE_STRUCTURED

  bit_length = writer.bit_length
  data = writer.to_bytes()
  payload = encode_int(bytes_to_int(data), alphabet)
  prefix = CODE_PREFIX_QR if qr else CODE_PREFIX_PAYLOAD
  return EncodeResult(
    code=prefix + payload,
    original_length=len(normalized),
    code_length=len(payload),
    bit_length=bit_length,
    mode=mode,
  )

def decode_url(code: str) -> str:
  if len(code) < 2:
    raise ValueError("code is too short")
  prefix, payload = code[0], code[1:]
  if prefix == CODE_PREFIX_QR:
    alphabet = QR_ALPHABET
  elif prefix == CODE_PREFIX_PAYLOAD:
    alphabet = PAYLOAD_ALPHABET
  else:
    raise ValueError("unknown code prefix")
  value = decode_int(payload, alphabet)
  byte_length = max(1, (value.bit_length() + 7) // 8)
  data = int_to_bytes(value, byte_length)

  # We don't know the exact bit length from base conversion, so parse with all possible leading zero counts.
  for leading_zero_bits in range(8):
    candidate = (b"\x00" * (1 if leading_zero_bits else 0)) + data
    bit_length = len(candidate) * 8 - leading_zero_bits
    try:
      return _decode_from_bytes(candidate, bit_length)
    except ValueError:
      continue
  raise ValueError("invalid or unsupported code")

def _decode_from_bytes(data: bytes, bit_length: int) -> str:
  reader = BitReader(data, bit_length)
  version = reader.read(4)
  if version not in SUPPORTED_CODE_VERSIONS:
    raise ValueError("unsupported version")
  qr_mode = reader.read(1)
  if qr_mode:
    return _read_raw(reader)
  return _read_structured_url(reader, version)

def _normalize_url(url: str) -> str:
  candidate = url.strip()
  if not candidate:
    raise ValueError("URL is empty")
  if "://" not in candidate:
    candidate = "https://" + candidate
  parts = urlsplit(candidate)
  if parts.scheme not in SCHEME_CODES:
    raise ValueError("only http and https URLs are supported")
  if not parts.netloc:
    raise ValueError("URL requires a host")
  return urlunsplit((parts.scheme.lower(), parts.netloc.lower(), parts.path, parts.query, parts.fragment))

def _write_structured_url(writer: BitWriter, url: str) -> None:
  parts = urlsplit(url)
  writer.write(SCHEME_CODES[parts.scheme], 1)

  host = parts.hostname or ""
  has_www = host.startswith("www.")
  if has_www:
    host = host[4:]
  writer.write(1 if has_www else 0, 1)

  labels = host.split(".") if host else []
  tld = labels[-1] if labels else ""
  sld = labels[-2] if len(labels) >= 2 else ""
  rest = ".".join(labels[:-2]) if len(labels) > 2 else ""

  _write_common_or_segment(writer, tld, COMMON_TLDS, TLD_CODE_WIDTH)
  _write_common_or_segment(writer, sld, COMMON_DOMAINS, DOMAIN_CODE_WIDTH)
  _write_text_segment(writer, rest)

  port = parts.port or 0
  writer.write(1 if port else 0, 1)
  if port:
    writer.write(port, 16)

  _write_text_segment(writer, parts.path)
  _write_text_segment(writer, parts.query)
  _write_text_segment(writer, parts.fragment)

def _read_structured_url(reader: BitReader, version: int) -> str:
  scheme = CODE_SCHEMES[reader.read(1)]
  has_www = bool(reader.read(1))
  dictionary_width = LEGACY_DICTIONARY_CODE_WIDTH if version == 1 else TLD_CODE_WIDTH
  tld = _read_common_or_segment(reader, COMMON_TLDS, dictionary_width)
  sld = _read_common_or_segment(reader, COMMON_DOMAINS, dictionary_width)
  rest = _read_text_segment(reader)

  labels = [part for part in [rest, sld] if part]
  host = ".".join(labels + ([tld] if tld else []))
  if has_www:
    host = "www." + host

  port = reader.read(16) if reader.read(1) else 0
  netloc = f"{host}:{port}" if port else host
  path = _read_text_segment(reader)
  query = _read_text_segment(reader)
  fragment = _read_text_segment(reader)
  return urlunsplit((scheme, netloc, path, query, fragment))

def _write_common_or_segment(writer: BitWriter, text: str, common: list[str], width: int) -> None:
  if text in common[1:]:
    writer.write(common.index(text), width)
  else:
    writer.write(0, width)
    _write_text_segment(writer, text)

def _read_common_or_segment(reader: BitReader, common: list[str], width: int) -> str:
  code = reader.read(width)
  if code:
    if code >= len(common):
      raise ValueError("invalid common dictionary code")
    return common[code]
  return _read_text_segment(reader)

def _write_text_segment(writer: BitWriter, text: str) -> None:
  if text == "":
    writer.write(0, 1)
    return
  writer.write(1, 1)

  compressed = _compress_text(text)
  raw = text.encode()
  use_compressed = len(compressed) + 1 < len(raw)
  writer.write(1 if use_compressed else 0, 1)
  data = compressed if use_compressed else raw
  writer.write(len(data), 16)
  for byte in data:
    writer.write(byte, 8)

def _read_text_segment(reader: BitReader) -> str:
  if not reader.read(1):
    return ""
  compressed = bool(reader.read(1))
  length = reader.read(16)
  data = bytes(reader.read(8) for _ in range(length))
  if compressed:
    data = zlib.decompress(data)
  return data.decode()

def _compress_text(text: str) -> bytes:
  return zlib.compress(text.encode(), level=9)

def _write_raw(writer: BitWriter, text: str) -> None:
  data = text.encode()
  writer.write(len(data), 16)
  for byte in data:
    writer.write(byte, 8)

def _read_raw(reader: BitReader) -> str:
  length = reader.read(16)
  return bytes(reader.read(8) for _ in range(length)).decode()

def describe_segments(url: str) -> list[dict[str, str | int]]:
  normalized = _normalize_url(url)
  chunks = [chunk for chunk in re.split(f"([{re.escape(SEPARATORS)}])", normalized) if chunk]
  return [_describe_segment(chunk) for chunk in chunks]

def _describe_segment(text: str) -> dict[str, str | int]:
  for key, alphabet in SEGMENT_ALPHABETS.items():
    if alphabet is not None and all(char in alphabet for char in text):
      return {"text": text, "alphabet": ALPHABET_NAMES[key], "size": len(alphabet)}
  return {"text": text, "alphabet": ALPHABET_NAMES[7], "size": 256}