"""Estrazione dell'anteprima incorporata nei file nativi Creo.

Creo salva nel file una miniatura JPEG dell'ultimo salvataggio: il render
ombreggiato del pezzo per i modelli, l'immagine della tavola per i disegni.
E' la stessa immagine che compare nella finestra "Apri" di Creo.

La miniatura non e' referenziata dall'indice UGC_TOC, quindi la si individua
per firma. Per non scambiare dati binari qualsiasi per un'immagine, ogni
candidato viene validato percorrendone i segmenti JPEG.
"""

from __future__ import annotations

import struct
from dataclasses import dataclass

SOI = b"\xff\xd8\xff"
EOI = b"\xff\xd9"

#: Marcatori "Start of Frame": contengono le dimensioni dell'immagine.
_SOF = {0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF}
_SOS = 0xDA
_STANDALONE = {0xD8, 0xD9, 0x01} | set(range(0xD0, 0xD8))

MIN_BYTES = 512


@dataclass
class Preview:
    """Una miniatura trovata nel file."""

    data: bytes
    offset: int
    width: int
    height: int
    format: str = "jpeg"

    def to_dict(self) -> dict:
        return {
            "present": True,
            "format": self.format,
            "width": self.width,
            "height": self.height,
            "bytes": len(self.data),
            "offset": hex(self.offset),
        }

    def save(self, path: str) -> None:
        with open(path, "wb") as handle:
            handle.write(self.data)


def _parse_jpeg(data: bytes, start: int) -> Preview | None:
    """Percorre i segmenti a partire da `start`. Restituisce None se la
    sequenza non e' un JPEG completo e coerente."""
    pos = start + 2  # oltre il SOI
    size = None

    while pos < len(data) - 3:
        if data[pos] != 0xFF:
            return None
        marker = data[pos + 1]

        if marker == 0xFF:      # byte di riempimento
            pos += 1
            continue
        if marker in _STANDALONE:
            pos += 2
            continue

        length = struct.unpack(">H", data[pos + 2 : pos + 4])[0]
        if length < 2:
            return None

        if marker in _SOF and pos + 9 <= len(data):
            height, width = struct.unpack(">HH", data[pos + 5 : pos + 9])
            size = (width, height)

        if marker == _SOS:
            # I dati compressi seguono il segmento: il byte FF che vi compare
            # e' sempre preceduto da escape, quindi il primo EOI e' la fine.
            end = data.find(EOI, pos + 2 + length)
            if end < 0 or size is None:
                return None
            blob = data[start : end + 2]
            if len(blob) < MIN_BYTES:
                return None
            return Preview(blob, start, size[0], size[1])

        pos += 2 + length

    return None


def find_all(data: bytes) -> "list[Preview]":
    """Tutte le miniature valide presenti nel file, in ordine di comparsa."""
    found, pos = [], 0
    while True:
        pos = data.find(SOI, pos)
        if pos < 0:
            return found
        preview = _parse_jpeg(data, pos)
        if preview:
            found.append(preview)
            pos += len(preview.data)
        else:
            pos += 1


def find(data: bytes) -> Preview | None:
    """La miniatura piu' grande, che e' quella a risoluzione migliore."""
    candidates = find_all(data)
    if not candidates:
        return None
    return max(candidates, key=lambda p: (p.width * p.height, len(p.data)))