#!/usr/bin/env python3
# TreeXcode — reference encoder
# Copyright 2026 Aleksandr Sitar <support@treexcode.com>
# SPDX-License-Identifier: Apache-2.0
# Licensed under the Apache License, Version 2.0 (see the LICENSE file or
# http://www.apache.org/licenses/LICENSE-2.0). Distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND.
"""
TreeXcode reference encoder.

Usage:
    python treexcode_encode.py "https://treexcode.com" -o code.png
    python treexcode_encode.py "hello world" --type text

Bakes in the standard interface parameters:
    version 0, medium error correction (Reed-Solomon, 16 parity bytes),
    automatic encoding (numeric / alphanumeric / UTF-8), jagged crown, diamond frame.
The only runtime input is the text; the output is a PNG image.

Format (see https://treexcode.com/math.html):
    byte0 = version(2) | type(4) | compressed(1) | encrypted(1)
    byte1 = encoding(3) | multipart(1) | ecc-level(4)
    length(2 bytes), then K Reed-Solomon blocks, each [chunk + CRC32(4)] + parity(nsym),
    chunk <= 251-nsym bytes, K = ceil(packedLen/(251-nsym)); multipart bit set when K>1.
    Capacity grows with tree size (more cells -> more blocks).

Dependencies: pillow.  (numpy optional, not required)
"""
import sys, argparse, math, zlib, random, re

# ---------------------------------------------------------------- glyph alphabet
DIR = {'N':(5,0),'S':(5,10),'E':(10,5),'W':(0,5),'NE':(10,0),'NW':(0,0),'SE':(10,10),'SW':(0,10),'C':(5,5)}
TILES = [([],[]),(['N','S'],[]),(['E','W'],[]),(['NE','SW'],[]),(['NW','SE'],[]),
 (['S','NE'],[]),(['S','NW'],[]),(['N','SE'],[]),(['N','SW'],[]),(['W','NE'],[]),(['W','SE'],[]),
 (['N','S','NE'],[]),(['N','S','NW'],[]),(['S','NE','NW'],[]),(['N','SE','SW'],[]),
 (['S'],['C']),(['SW'],['C']),(['SE'],['C']),([],['C']),(['E','NW'],[]),(['E','SW'],[]),
 (['NW'],['C']),(['NE'],['C']),(['N'],['C']),(['W'],['C']),(['E'],['C']),None]
GREEN = 26
B36 = '0123456789abcdefghijklmnopqrstuvwxyz'
SIM = {'ul':[2,4,17,18,25,0],'ur':[2,3,16,18,24,0],'ll':[2,3,22,18,25,0],'lr':[2,4,21,18,24,0]}
# Data uses 4 of the 6 glyph shapes -> 2 bits per cell, 4 cells per byte, so one misread
# cell equals at most one wrong byte and Reed-Solomon stays effective. The centre dot and the
# empty cell are part of the alphabet but are not used for payload.
TYPES4 = ['full_h','full_d','half_h','half_d']
def tile_type(t):
    if t==0:return'empty'
    if t==18:return'dot'
    if t==2:return'full_h'
    if t in(3,4):return'full_d'
    if t in(24,25):return'half_h'
    return'half_d'
Q2T = {q:{TYPES4.index(tile_type(t)):t for t in ts if tile_type(t) in TYPES4} for q,ts in SIM.items()}

ORN_ENC = '37/36f/35iA/33f1fA/33AfAA/33AAAA/33GGAA/34FAA/33fAnA/33AGiA/33G1FA/34FAA/33iAAA/34AAA/33DnnA/34F1A/34A1A/33DAfA/33FnAA/33niAA/33DfAA/34AAA/33lfAA/33DAAA/34AAA/33DfAA/33DAAA/34DAA/33DiAA/34FAn/33lAAf/34nnA/33Df1A/3kBBBBBBBI1kBBI1pBokBBBa1pBIC1AaA1A/2m1kBBI4m2kBBBokBBapIDi1C2GBAfA/3m10m4C4jBBoC4DnAA/18C8m6FAA/17m16AAA/34AAA/34AAA/34AAA/34AAA/34A2/37/37/37/34A2/34AAA/34AAA/34AAA/34AnA/33CAfA/29pBBICHAA/5pBBBBBBBoppBBBBBBBoipBBBBBBBI1AA/6pBBBBBBogCmCpIkBBBBBBBopBBoCAAA/8kBBBBIm1C2C1kBo1kBBBBBBI1AAA/7m25CA1A/34AAA/34HAA/33CfAA/33fAAA/33A1AA/33A1HA/33nE1A/33iAfA/34AAA/34AAA/34AnA/34AiA/34n1A/36A/36i/37'
def _decode_orn(s):
    out=[]
    for rs in s.split('/'):
        row=''; i=0
        while i<len(rs):
            ch=rs[i]
            if ch.isdigit():
                j=i
                while j<len(rs) and rs[j].isdigit(): j+=1
                row+='.'*int(rs[i:j]); i=j
            else:
                row+=(str(ord(ch)-64) if 'A'<=ch<='I' else ch); i+=1
        out.append(row)
    return out
ORN = _decode_orn(ORN_ENC)

def node_rcells(N): return 0.5 + 0.25*max(0, N//32 - 1)

def gen_tree(N, edge='jagged'):
    grid=[[0]*N for _ in range(N)]
    c=N/2; Rz=0.47*N; ht=max(1,round(N*0.045))
    crownB=round(32*(N-1)/72); rootT=round(56*(N-1)/72); nodeY=round(N*0.6)
    STEP=max(2,round(N*0.04))
    def in_disk(x,y):
        dx=x+0.5-c; dy=y+0.5-c
        if edge=='jagged':
            by=math.ceil(abs(dy)/STEP)*STEP; xm2=Rz*Rz-by*by
            return xm2>0 and dx*dx<xm2
        return dx*dx+dy*dy<Rz*Rz
    for y in range(N):
        for x in range(N):
            if in_disk(x,y) and abs(x+0.5-c)>ht and (y<=crownB or y>=rootT): grid[y][x]=GREEN
    rows=ORN; W1=len(rows)-1
    MIR={3:4,4:3,5:6,6:5,7:8,8:7,9:19,19:9,10:20,20:10,16:17,17:16,21:22,22:21,24:25,25:24}
    half=N//2
    for y in range(N):
        src=rows[round(y*W1/(N-1))]
        for x in range(half):
            v=B36.find(src[min(36,round(x*W1/(N-1)))])
            if v>0 and v!=GREEN and grid[y][x]!=GREEN: grid[y][x]=v
            xm=N-1-x; vm=MIR.get(v,v) if v>0 else 0
            if vm and grid[y][xm]!=GREEN: grid[y][xm]=vm
        if N%2==1:
            v=B36.find(src[min(36,round(half*W1/(N-1)))])
            if v>0 and v!=GREEN and grid[y][half]!=GREEN: grid[y][half]=v
    rc=2*math.sqrt(2)*node_rcells(N)+0.5
    for y in range(N):
        for x in range(N):
            if (x+0.5-c)**2+(y-nodeY)**2 < rc*rc: grid[y][x]=0
    for rx,ry in [(N/2,2),(N/2,N-2),(2,N/2),(N-2,N/2)]:
        for y in range(N):
            for x in range(N):
                if (x+0.5-rx)**2+(y+0.5-ry)**2 < 1.6*1.6: grid[y][x]=0
    greens=[(x,y) for y in range(N) for x in range(N) if grid[y][x]==GREEN]
    return grid, greens, nodeY

def quad_of(x,y,N):
    up = y < round(N*0.6); left = x < N//2
    return ('ul' if left else 'ur') if up else ('ll' if left else 'lr')

# ---------------------------------------------------------------- Reed-Solomon GF(256)
EXP=[0]*512; LOG=[0]*256
_x=1
for _i in range(255):
    EXP[_i]=_x; LOG[_x]=_i; _x<<=1
    if _x&0x100: _x^=0x11d
for _i in range(255,512): EXP[_i]=EXP[_i-255]
def _gmul(a,b): return 0 if a==0 or b==0 else EXP[LOG[a]+LOG[b]]
def _gpow(a,p): return EXP[(LOG[a]*p)%255]
def _pmul(p,q):
    r=[0]*(len(p)+len(q)-1)
    for j in range(len(q)):
        for i in range(len(p)): r[i+j]^=_gmul(p[i],q[j])
    return r
def _rs_gen(nsym):
    g=[1]
    for i in range(nsym): g=_pmul(g,[1,_gpow(2,i)])
    return g
def rs_encode(msg,nsym):
    g=_rs_gen(nsym); res=list(msg)+[0]*(len(g)-1)
    for i in range(len(msg)):
        c=res[i]
        if c:
            for j in range(1,len(g)): res[i+j]^=_gmul(g[j],c)
    return list(msg)+res[len(msg):]
ECC_NSYM=[0,8,16,32]

# ---------------------------------------------------------------- encodings
AN='0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ $%*+-./:'
def _is_digits(s): return bool(s) and all('0'<=ch<='9' for ch in s)
def _is_alnum(s): return bool(s) and all(ch in AN for ch in s)
def detect_enc(text):
    if _is_digits(text): return 1
    if _is_alnum(text): return 2
    return 0
def _bits_to_bytes(bits):
    out=[]
    for i in range(0,len(bits),8):
        b=0
        for j in range(8): b=(b<<1)|(bits[i+j] if i+j<len(bits) else 0)
        out.append(b)
    return out
def _push(bits,v,n):
    for i in range(n-1,-1,-1): bits.append((v>>i)&1)
def pack_content(text,mode):
    if mode==1:
        bits=[]; i=0
        while i+3<=len(text): _push(bits,int(text[i:i+3]),10); i+=3
        r=len(text)-i
        if r==2: _push(bits,int(text[i:i+2]),7)
        elif r==1: _push(bits,int(text[i:i+1]),4)
        return _bits_to_bytes(bits)
    if mode==2:
        bits=[]; i=0
        while i+2<=len(text): _push(bits,AN.index(text[i])*45+AN.index(text[i+1]),11); i+=2
        if len(text)-i==1: _push(bits,AN.index(text[i]),6)
        return _bits_to_bytes(bits)
    if mode==3: return [ord(c)&0xff for c in text]
    return list(text.encode('utf-8'))
def logical_len(text,mode): return len(text.encode('utf-8')) if mode==0 else len(text)

# ---------------------------------------------------------------- content type
def auto_type(text):
    t=text.strip()
    if re.match(r'^[^\s@]+@[^\s@]+\.[^\s@]+$', t): return 3          # email
    if re.match(r'^(https?://|www\.)\S+$', t, re.I) or re.match(r'^[a-z0-9.-]+\.[a-z]{2,}(/\S*)?$', t, re.I): return 1
    if re.match(r'^[+]?[\d\s().-]{6,}$', t): return 4                # phone
    return 0

# ---------------------------------------------------------------- capacity / size
N_LADDER=[48,56,64,72,80,88,96,104,112,120]     # minimum code is 48x48
def capacity_bytes(N, level):
    # multipart: заголовок 4 байта + K RS-блоков по 255 байт (D=251-nsym данных + 4 CRC + nsym parity)
    _,greens,_=gen_tree(N); nsym=ECC_NSYM[level]
    R=len(greens)//4 - 4
    if R<=0: return 0
    D=max(1,251-nsym); nfull=R//255; rem=R-nfull*255
    return min(nfull*D + max(0, rem-4-nsym), 65535)
def auto_size(nbytes, level):
    for N in N_LADDER:
        if capacity_bytes(N,level)>=nbytes: return N
    raise ValueError('content too large for the largest supported code')

# ---------------------------------------------------------------- encode
def encode(text, type=None, ecc=2, seed=None):
    typ = auto_type(text) if type is None else type
    mode = detect_enc(text)
    content = pack_content(text, mode); L = logical_len(text, mode)
    N = auto_size(len(content), ecc)
    grid, greens, nodeY = gen_tree(N)
    nsym=ECC_NSYM[ecc]; PB=len(content); D=max(1,251-nsym)
    K=max(1, math.ceil(PB/D)); multipart=1 if K>1 else 0
    byte0 = (0<<6)|((typ&15)<<2); byte1 = ((mode&7)<<5)|((multipart&1)<<4)|(ecc&15)
    full = [byte0,byte1,(L>>8)&255,L&255]
    for b in range(K):                                   # K независимых RS-блоков
        chunk = content[b*D: min((b+1)*D, PB)]
        crc = zlib.crc32(bytes(chunk)) & 0xffffffff
        block = chunk + [(crc>>24)&255,(crc>>16)&255,(crc>>8)&255,crc&255]
        if nsym>0: block = rs_encode(block, nsym)
        full += block
    syms=[]
    for b in full: syms += [(b>>6)&3,(b>>4)&3,(b>>2)&3,b&3]
    rnd = random.Random(seed if seed is not None else 0x51ab13)
    greens.sort(key=lambda p:(p[1],p[0]))
    for k,(x,y) in enumerate(greens):
        v = syms[k] if k<len(syms) else rnd.randrange(4)
        grid[y][x] = Q2T[quad_of(x,y,N)][v]
    crc_all = zlib.crc32(bytes(content)) & 0xffffffff
    return grid, N, nodeY, dict(type=typ, mode=mode, ecc=ecc, length=L, crc=crc_all, used=len(syms), cells=len(greens), blocks=K)

# ---------------------------------------------------------------- render (PIL)
def render(grid, N, nodeY, cs=20, fg=(20,35,27), bg=(255,255,255), leaf=(42,138,63)):
    from PIL import Image, ImageDraw
    U=N*cs; img=Image.new('RGB',(U,U),bg); dr=ImageDraw.Draw(img)
    lw=max(2,round(cs*0.15)); s=cs/10.0
    for y in range(N):
        for x in range(N):
            t=grid[y][x]
            if not t: continue
            ox=x*cs; oy=y*cs
            if t==GREEN:
                dr.rectangle([ox+cs*0.06,oy+cs*0.06,ox+cs*0.94,oy+cs*0.94], fill=leaf); continue
            spk,dts=TILES[t]; cc=(ox+5*s,oy+5*s)
            for d in spk:
                p=(ox+DIR[d][0]*s, oy+DIR[d][1]*s); dr.line([cc,p], fill=fg, width=lw)
            for d in dts:
                p=(ox+DIR[d][0]*s, oy+DIR[d][1]*s); rr=cs*0.17; dr.ellipse([p[0]-rr,p[1]-rr,p[0]+rr,p[1]+rr], fill=fg)
    cx=N/2*cs; cy=(round(N*0.6)+0.5)*cs; r=node_rcells(N)*cs; D=2*math.sqrt(2)*r
    dr.line([(cx,cy-D),(cx+D,cy),(cx,cy+D),(cx-D,cy),(cx,cy-D)], fill=fg, width=lw, joint='curve')
    off=0.5*r/math.sqrt(2); mx=cx+D/2+off; my=cy-D/2-off; hl=(cs/2)/math.sqrt(2)
    dr.line([(mx-hl,my-hl),(mx+hl,my+hl)], fill=fg, width=lw)
    dr.ellipse([cx-r,cy-r,cx+r,cy+r], fill=fg)
    ca=r/2; dd=ca*1.4
    for px,py in [(N/2*cs,2*cs),(N/2*cs,U-2*cs),(2*cs,N/2*cs),(U-2*cs,N/2*cs)]:
        dr.polygon([(px,py-dd),(px+dd,py),(px,py+dd),(px-dd,py)], fill=fg)
    return img

# ---------------------------------------------------------------- CLI
def main():
    ap=argparse.ArgumentParser(description='Encode text into a TreeXcode image.')
    ap.add_argument('text', help='content to encode')
    ap.add_argument('-o','--out', default='treexcode.png', help='output PNG path')
    ap.add_argument('-t','--type', default='auto',
                    help='content type: auto|text|url|email|phone|sms|geo|wifi|vcard|event')
    ap.add_argument('--cell', type=int, default=20, help='pixels per cell')
    a=ap.parse_args()
    names=['text','url','treex-short','email','phone','sms','geo','wifi','vcard','event','payment','binary']
    typ = None if a.type=='auto' else names.index(a.type)
    grid,N,nodeY,meta = encode(a.text, type=typ, ecc=2)
    img = render(grid, N, nodeY, cs=a.cell)
    img.save(a.out)
    print('Saved %s  (%dx%d tree, type=%s, encoding=%s, ecc=medium, %d/%d symbols)'
          % (a.out, N, N, names[meta['type']], ['UTF-8','Numeric','Alphanumeric','Latin-1'][meta['mode']], meta['used'], meta['cells']))

if __name__=='__main__':
    main()