__all__ = [
    'floor2',
    'gdal_parse_mpp',
    'get_aperio_properties',
    'get_ventana_properties',
    'parse_xml',
    'round2',
    'unflatten',
]

from collections.abc import Mapping
from io import BytesIO
from itertools import zip_longest
from typing import Any

import cv2
import numpy as np
from lxml import etree
from lxml.etree import XMLParser, fromstring
from PIL.Image import fromarray
from PIL.ImageCms import buildTransform, createProfile

# ------------------------- flat list of properties --------------------------


def unflatten(d: Mapping[str, str]) -> dict[str, Any]:
    if not d:
        return {}

    tab = {ord('['): '.', ord(']'): None}
    pairs = sorted((_make_key(k, tab), v) for k, v in d.items())
    r: dict = {}

    for parts, v in pairs:
        rr = r
        for curr, next_ in zip_longest(parts, parts[1:], fillvalue=None):
            # Put list if next is list idx, otherwise put dict
            default: list | dict | str
            default = (
                v if next_ is None else ([] if isinstance(next_, int) else {})
            )

            if isinstance(curr, int):  # last was list idx
                if not isinstance(rr, list):  # Ambigous structure
                    break
                if curr >= len(rr):
                    rr.extend(None for _ in range(len(rr), curr + 1))
                    rr[curr] = default
                rr = rr[curr]
            else:  # last was dict key
                if not isinstance(rr, dict):  # Ambigous structure
                    break
                rr = rr.setdefault(curr, default)
    return r


def _make_key(s: str, tab: Mapping[int, str | None]) -> tuple[str | int, ...]:
    parts = s.translate(tab).split('.')
    return tuple(int(word) if word.isdigit() else word for word in parts)


# ---------------------------- aperio description ----------------------------


def get_aperio_properties(
    description: str, index: int = 0
) -> tuple[str, dict[str, str]] | None:
    if index == 0 and not description.startswith('Aperio'):
        return None
    header, *kv_pairs = description.split('|')

    header = '\n'.join(s.strip() for s in header.splitlines())

    meta = {}
    for i, kv in enumerate(kv_pairs):
        match kv.split('=', 1):
            case [str(k), str(v)]:
                meta[k.strip()] = v.strip()
            case ['']:
                continue
            case _:
                raise ValueError(
                    f'Cannot parse TIFF description line #{i}: '
                    f'{kv!r}, {description!r}'
                )

    return header, meta


# ----------------------------- xml description ------------------------------


def get_ventana_properties(s: bytes, index: int = 0) -> dict[str, str]:
    s = s.strip(b'\00')  # For safety
    t = fromstring(s, XMLParser(resolve_entities=False, no_network=True))
    if index == 0:
        if (root := t.find('iScan')) is not None:
            return dict(root.items())
        return {}
    return {'xmp': s.decode()}


def parse_xml(
    s: str,
    /,
    *,
    group: str = 'property',
    name: str = 'name',
    value: str = 'value',
) -> dict[str, str]:
    t = fromstring(s, XMLParser(resolve_entities=False, no_network=True))

    # Remove a namespace URI in the element's name
    for elem in t.getiterator():
        if not isinstance(elem, etree._Comment | etree._ProcessingInstruction):
            elem.tag = etree.QName(elem).localname

    # Remove unused namespace declarations
    etree.cleanup_namespaces(t, top_nsmap=None, keep_ns_prefixes=None)

    return {e.find(name).text: e.find(value).text for e in t.iter(group)}


# ----------------------------------------------------------------------------


def _gdal_parse_description(meta: Mapping[str, str]) -> dict[str, Any]:
    desc = meta.get('TIFFTAG_IMAGEDESCRIPTION')
    if not desc:
        return {}
    if get_aperio_properties(desc):
        raise ValueError('Aperio is not yet supported by GDAL driver')
    try:
        props = parse_xml(desc)
    except Exception:  # noqa: BLE001
        return {}
    return unflatten(props)


def gdal_parse_mpp(meta: Mapping) -> list[float]:
    props = _gdal_parse_description(meta)

    # TIFF tags
    res = [float(v) for a in 'XY' if (v := meta.get(f'TIFFTAG_{a}RESOLUTION'))]
    if res:
        match meta.get('TIFFTAG_RESOLUTIONUNIT'):
            case '2 (pixels/inch)':
                return [25_400 / r for r in res]
            case '3 (pixels/cm)':
                return [10_000 / r for r in res]

    # VIPS tags, always px/mm
    if res := [float(v) for tag in ('xres', 'yres') if (v := props.get(tag))]:
        return [1_000 / r for r in res]

    # Openslide tags
    if osd := props.get('openslide'):
        mpp = [float(v) for tag in ('mpp-x', 'mpp-y') if (v := osd.get(tag))]
        if mpp:
            return mpp

    # Aperio tags
    if mpp_ := props.get('MPP'):
        return [float(mpp_), float(mpp_)]

    return []


# ------------- contrast-limited adaptive histogram equalization -------------


def clahe(
    im: np.ndarray, cell_size: int = 64, clip_limit: float = 2.0
) -> np.ndarray:
    h, w = im.shape[:2]
    gh, gw = (h // cell_size) or 1, (w // cell_size) or 1
    cl = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(gw, gh))

    ls, a, b = cv2.split(cv2.cvtColor(im, cv2.COLOR_RGB2LAB))
    ls = cl.apply(ls)
    return cv2.cvtColor(cv2.merge([ls, a, b]), cv2.COLOR_LAB2RGB)


# -------------------------- image color correction --------------------------


class Icc:
    def __init__(self, icc: bytes) -> None:
        f = BytesIO(icc)
        srgb = createProfile('sRGB')
        self._tf = buildTransform(f, srgb, inMode='RGB', outMode='RGB')

    def __call__(self, image: np.ndarray) -> np.ndarray:
        pil = fromarray(image)
        pil = self._tf.apply(pil)
        return np.array(pil, copy=False)


# -------------------------------- etc---------------------------------------


def round2(x: float) -> int:
    """Round to `2^k`, `k` is non-negative"""
    return floor2(x * _sqrt2)


def floor2(x: float) -> int:
    """Floor to `2^k`, `k` is non-negative"""
    x = max(x, 1)
    power = int(x).bit_length() - 1
    return 1 << power


_sqrt2 = 2**0.5
