# Part of the Hex9 (H9) Project
# Copyright ©2025, Ben Griffin
# Licensed under the Apache License, Version 2.0

"""
H9 UUID Address Encoding.

Provides a stable 128-bit (UUID) representation of a hex9 address at maximum
practical depth, plus a companion byte that enables exact round-trip to lat/lon.

Layout (32 nibbles = 128 bits):
    nibbles  0..30 : hex body (L0..L30) — hierarchical hex digits
    nibble      31 : key tail = (p_c2 << 1) | r_mo   (values 0..5)

Companion adr byte (1 byte, 5 bits used):
    bit 4      : p_mo  — parent mode of terminal region
    bits 3..0  : h     — terminal region id (0..11)

The UUID alone is sufficient for spatial indexing and hexbin at any layer.
The (UUID, adr) pair enables exact reconstruction of the b_oct coordinate
and therefore round-trip to lat/lon.

Public API:
    h9_encode(lats, lons)              -> (uuids, adr_bytes)
    h9_decode(uuids, adr_bytes)        -> (lats, lons)
    h9_bin(uuids, layer)               -> uuids
"""

from __future__ import annotations

import uuid as uuid_mod
import numpy as np
from numpy.typing import NDArray

from hhg9.h9.addressing import hex_digits, hex_decode, H9_RA
from hhg9.h9.tail import TailStyle, tail_pack_reversible, tail_unpack_reversible
from hhg9.h9.protocols import RegionAddressLike

# The layer parameter passed to hex_digits that yields 32 nibbles:
#   xy_regions(depth=30) -> addresses shape (N, 32)
#   reg_hex_digits sees cols=32, depth=31 -> bdy has 31 cols (L0..L30)
#   key pack: 31 body nibbles + 1 key_tail nibble = 32 nibbles = 128 bits
UUID_DEPTH: int = 30


# ---------- 128-bit packing helpers ----------------------------------------

def _nibbles_to_int(nibbles: NDArray[np.uint8]) -> int:
    """Pack a (32,) uint8 nibble row into a Python int, MSB-first."""
    hi = nibbles[0::2].astype(np.uint8)
    lo = nibbles[1::2].astype(np.uint8)
    return int.from_bytes(bytes((hi << 4) | lo), 'big')


def _batch_nibbles_to_int(nibbles: NDArray[np.uint8]) -> list[int]:
    """Pack (N, 32) uint8 nibbles into N Python ints, MSB-first."""
    hi = nibbles[:, 0::2].astype(np.uint8)
    lo = nibbles[:, 1::2].astype(np.uint8)
    byte_arr = (hi << 4) | lo  # (N, 16)
    return [int.from_bytes(bytes(row), 'big') for row in byte_arr]


def _batch_int_to_nibbles(values: list[int], n: int = 32) -> NDArray[np.uint8]:
    """Unpack N Python ints into (N, n) uint8 nibbles, MSB-first."""
    n_bytes = n // 2  # n is always even (32)
    result = np.zeros((len(values), n), dtype=np.uint8)
    for i, val in enumerate(values):
        b = val.to_bytes(16, 'big')[-n_bytes:]  # rightmost n_bytes of 128-bit int
        for j, byte in enumerate(b):
            result[i, 2 * j] = (byte >> 4) & 0xF
            result[i, 2 * j + 1] = byte & 0xF
    return result


# ---------- Core encode/decode ---------------------------------------------

def h9_encode(
        lats,
        lons,
        reg=None,
        scheme: RegionAddressLike = H9_RA,
) -> tuple[list[uuid_mod.UUID], NDArray[np.uint8]]:
    """
    Encode geographic coordinates to H9 UUID addresses.

    Parameters
    ----------
    lats : array-like, degrees, WGS84 geodetic latitude
    lons : array-like, degrees, WGS84 geodetic longitude
    reg  : Registrar (created if None)
    scheme : RegionAddressLike (normally H9_RA)

    Returns
    -------
    uuids : list[uuid.UUID]  — 128-bit key addresses, one per point
    adr_bytes : NDArray[uint8] — companion byte enabling round-trip decode,
                packed as (p_mo << 4) | h
    """
    from hhg9 import Registrar, Points

    lats = np.asarray(lats, dtype=np.float64).ravel()
    lons = np.asarray(lons, dtype=np.float64).ravel()
    if lats.shape != lons.shape:
        raise ValueError("lats and lons must have the same length")

    if reg is None:
        reg = Registrar()

    g_gcd = reg.domain('g_gcd')
    b_oct = reg.domain('b_oct')

    # g_gcd uses [lat, lon] column order
    coords = np.column_stack([lats, lons])
    b_pts = reg.project(Points(coords, g_gcd), [g_gcd, b_oct])

    # Full reversible address at UUID_DEPTH
    hx = hex_digits(b_pts, layer=UUID_DEPTH, tail_style=TailStyle.reversible, scheme=scheme)
    # hx shape: (N, 32) — 31 body nibble cols + 1 full tail byte col

    body = hx[:, :-1]          # (N, 31): L0..L30 as nibble values
    tail_bytes = hx[:, -1]     # (N,):    reversible tail byte

    # Key tail: high nibble of reversible tail = (p_c2 << 1) | r_mo
    key_tail_nibble = ((tail_bytes >> 4) & 0x0F).astype(np.uint8)  # values 0..5

    # UUID nibbles: 31 body + 1 key_tail = 32
    uuid_nibbles = np.column_stack([body, key_tail_nibble])  # (N, 32)

    # Adr byte: the two fields absent from the key tail
    p_mo = ((tail_bytes >> 7) & 0x01).astype(np.uint8)   # 1 bit
    h = (tail_bytes & 0x0F).astype(np.uint8)              # 4 bits  (region 0..11)
    adr_bytes = ((p_mo << 4) | h).astype(np.uint8)        # (N,)

    uuids = [uuid_mod.UUID(int=v) for v in _batch_nibbles_to_int(uuid_nibbles)]
    return uuids, adr_bytes


def h9_decode(
        uuids: list[uuid_mod.UUID],
        adr_bytes,
        reg=None,
        scheme: RegionAddressLike = H9_RA,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """
    Decode H9 (UUID, adr) pairs back to geographic coordinates.

    Parameters
    ----------
    uuids     : list[uuid.UUID] — key UUIDs from h9_encode
    adr_bytes : array-like uint8 — companion adr bytes from h9_encode
    reg       : Registrar (created if None)
    scheme    : RegionAddressLike (normally H9_RA)

    Returns
    -------
    lats : NDArray[float64], geodetic latitude  (degrees)
    lons : NDArray[float64], geodetic longitude (degrees)
    """
    from hhg9 import Registrar

    if reg is None:
        reg = Registrar()

    b_oct = reg.domain('b_oct')
    g_gcd = reg.domain('g_gcd')

    uuid_ints = [u.int for u in uuids]
    uuid_nibbles = _batch_int_to_nibbles(uuid_ints, n=32)   # (N, 32)
    adr_bytes = np.asarray(adr_bytes, dtype=np.uint8)

    # Recover the four fields that make the reversible tail
    key_tail_nibble = uuid_nibbles[:, 31]                        # (p_c2 << 1) | r_mo
    p_c2 = ((key_tail_nibble >> 1) & 0x3).astype(np.uint8)
    r_mo = (key_tail_nibble & 0x1).astype(np.uint8)
    p_mo = ((adr_bytes >> 4) & 0x1).astype(np.uint8)
    h = (adr_bytes & 0x0F).astype(np.uint8)

    reversible_tail = tail_pack_reversible(p_mo, p_c2, r_mo, h)  # (N,) uint8

    # Reconstruct the full hx array that hex_decode expects
    body = uuid_nibbles[:, :31]                            # (N, 31) nibble values
    hx = np.column_stack([body, reversible_tail])          # (N, 32)

    b_pts = hex_decode(hx, b_oct, scheme=scheme)
    g_pts = reg.project(b_pts, [b_oct, g_gcd])

    lats = g_pts.coords[:, 0]
    lons = g_pts.coords[:, 1]
    return lats, lons


def h9_bin(
        uuids: list[uuid_mod.UUID],
        layer: int,
        reg=None,
        scheme: RegionAddressLike = H9_RA,
) -> list[uuid_mod.UUID]:
    """
    Return the layer-L bin UUID for each H9 UUID address.

    The bin UUID has the same 128-bit format: the first (layer+2) nibbles carry
    the hierarchical bin key at the requested layer; the remaining nibbles are zero.
    Bin UUIDs at the same layer can be compared directly for equality and sorted.

    The adr companion byte is NOT required — the key UUID alone contains sufficient
    information to identify the correct bin at any layer <= UUID_DEPTH.

    Parameters
    ----------
    uuids : list[uuid.UUID] — key UUIDs from h9_encode
    layer : int — target hex layer (0..UUID_DEPTH)
    reg   : Registrar (created if None)

    Returns
    -------
    list[uuid.UUID] — bin-key UUIDs at the requested layer
    """
    from hhg9 import Registrar

    if not (0 <= layer <= UUID_DEPTH):
        raise ValueError(f"layer must be in 0..{UUID_DEPTH}, got {layer}")

    if reg is None:
        reg = Registrar()

    b_oct = reg.domain('b_oct')

    uuid_ints = [u.int for u in uuids]
    uuid_nibbles = _batch_int_to_nibbles(uuid_ints, n=32)   # (N, 32)

    # Reconstruct a key-style tail byte from the stored key nibble.
    # This gives an approximate b_oct point (p_mo forced to 0, terminal h=0xF)
    # that is guaranteed to lie within the correct depth-30 hexagon bin,
    # and therefore within the correct depth-L bin for any L <= UUID_DEPTH.
    key_tail_nibble = uuid_nibbles[:, 31]
    key_tail_byte = ((key_tail_nibble.astype(np.uint8) << 4) | np.uint8(0x0F))

    body = uuid_nibbles[:, :31]                            # (N, 31)
    hx_full = np.column_stack([body, key_tail_byte])       # (N, 32)

    approx_pts = hex_decode(hx_full, b_oct, scheme=scheme)

    # Recompute the address at the target layer with a key tail
    hx_L = hex_digits(approx_pts, layer=layer, tail_style=TailStyle.key, scheme=scheme)
    # hx_L shape: (N, layer+2) — (layer+1) body nibbles + 1 tail byte

    body_L = hx_L[:, :-1]        # (N, layer+1): L0..Llayer
    tail_byte_L = hx_L[:, -1]    # (N,) key-style tail byte
    key_nibble_L = ((tail_byte_L >> 4) & 0x0F).astype(np.uint8)

    # Pack into 32-nibble UUID: significant prefix + zero padding
    uuid_nibs_L = np.zeros((len(uuids), 32), dtype=np.uint8)
    uuid_nibs_L[:, :layer + 1] = body_L    # body at layer L
    uuid_nibs_L[:, layer + 1] = key_nibble_L  # key tail at layer L
    # nibbles[layer+2..31] remain zero

    return [uuid_mod.UUID(int=v) for v in _batch_nibbles_to_int(uuid_nibs_L)]
