"""True 16-bit RGB BT.2020/PQ PNG, using PNG 3 cICP and cLLI metadata.

Color math and compositing operate in absolute linear-light nits. resvg is used
only for vector coverage masks (8-bit alpha), never to supply HDR color values.
"""

from copy import deepcopy
import io
from pathlib import Path
import struct
import xml.etree.ElementTree as ET
import zlib

import numpy as np
from PIL import Image
import resvg_py

from avatar_color import BT2020_Y, equal_luminance, validate_peak

NS = '{http://www.w3.org/2000/svg}'


def pq_encode(nits):
    """ST 2084 inverse EOTF; 10,000 nits maps to 1."""
    m1, m2 = 2610 / 16384, 2523 / 32
    c1, c2, c3 = 3424 / 4096, 2413 / 128, 2392 / 128
    p = np.maximum(np.asarray(nits, dtype=np.float64), 0) / 10000
    p = p ** m1
    return ((c1 + c2 * p) / (1 + c3 * p)) ** m2


def pq_decode(signal):
    m1, m2 = 2610 / 16384, 2523 / 32
    c1, c2, c3 = 3424 / 4096, 2413 / 128, 2392 / 128
    p = np.asarray(signal, dtype=np.float64) ** (1 / m2)
    return 10000 * (np.maximum(p - c1, 0) / (c2 - c3 * p)) ** (1 / m1)


def render_linear_hdr(svg, peak_nits=400):
    validate_peak(peak_nits)
    root = ET.fromstring(svg)
    width, height = int(root.get('width')), int(root.get('height'))
    canvas = np.full((height, width, 3), .02, dtype=np.float32)
    # Match SVG paint order. The current exporter has a translated background
    # group and an independently transformed logo path.
    layers = []
    for child in root:
        if child.tag == NS + 'g':
            for path in child:
                layers.append((path, child.get('transform', '')))
        elif child.tag == NS + 'path':
            layers.append((child, ''))
    for path, parent_transform in layers:
        base = path.get('data-base')
        if base is None:
            raise ValueError('HDR rendering requires palette metadata from make_svg')
        level = float(path.get('data-level', '1'))
        peak = equal_luminance(base, peak_nits, 10000, hdr=True)
        rgb = np.asarray(peak, dtype=np.float32) * level
        mask_svg = ET.Element(NS + 'svg', dict(root.attrib))
        mask_path = deepcopy(path)
        mask_path.set('fill', '#ffffff')
        mask_path.set('transform', parent_transform + ' ' + path.get('transform', ''))
        mask_svg.append(mask_path)
        png = resvg_py.svg_to_bytes(svg_string=ET.tostring(mask_svg, encoding='unicode'))
        with Image.open(io.BytesIO(png)) as image:
            coverage = image.convert('RGBA').getchannel('A')
            box = coverage.getbbox()
            if box is None:
                continue
            alpha = np.asarray(coverage.crop(box), dtype=np.float32)[..., None] / 255
        x0, y0, x1, y1 = box
        region = canvas[y0:y1, x0:x1]
        region *= 1 - alpha
        region += alpha * rgb
    return canvas


def png_chunk(kind, payload):
    return (struct.pack('>I', len(payload)) + kind + payload
            + struct.pack('>I', zlib.crc32(kind + payload) & 0xffffffff))


def write_hdr_png(path, linear_nits):
    height, width, channels = linear_nits.shape
    if channels != 3 or not np.isfinite(linear_nits).all() or linear_nits.min() < 0 or linear_nits.max() > 10000:
        raise ValueError('HDR pixels must be finite RGB values in 0..10000 nits')
    encoded = np.rint(pq_encode(linear_nits) * 65535).astype('>u2')
    # Calculate static luminance metadata from the actually encoded pixels.
    decoded = pq_decode(encoded.astype(np.float64) / 65535)
    luminance = decoded @ np.asarray(BT2020_Y)
    max_y, average_y = float(luminance.max()), float(luminance.mean())
    scanlines = b''.join(b'\0' + row.tobytes() for row in encoded)
    content = b'\x89PNG\r\n\x1a\n'
    content += png_chunk(b'IHDR', struct.pack('>IIBBBBB', width, height, 16, 2, 0, 0, 0))
    content += png_chunk(b'cICP', bytes((9, 16, 0, 1)))  # BT.2020, PQ, RGB, full range
    content += png_chunk(b'cLLI', struct.pack('>II', round(max_y * 10000), round(average_y * 10000)))
    content += png_chunk(b'tEXt', b'Software\0cmatrix-avatar; linear-light BT.2020/PQ renderer')
    content += png_chunk(b'IDAT', zlib.compress(scanlines, 6))
    content += png_chunk(b'IEND', b'')
    Path(path).write_bytes(content)
    return max_y, average_y


def export_hdr(svg, output, peak_nits=400):
    return write_hdr_png(output, render_linear_hdr(svg, peak_nits))