pyheic_struct/heic_file.py

# pyheic_struct/heic_file.py

import os
import math
# 1. 导入新的库
import pillow_heif
from io import BytesIO
from dataclasses import dataclass
from PIL import Image

from base import Box
from parser import parse_boxes
from heic_types import (
    ItemLocationBox, PrimaryItemBox, ItemInfoBox, ItemPropertiesBox,
    ImageSpatialExtentsBox, ItemReferenceBox, ItemInfoEntryBox,
    _read_int
)
from handlers.base_handler import VendorHandler
from handlers.apple_handler import AppleHandler
from handlers.samsung_handler import SamsungHandler

@dataclass
class Grid:
    rows: int
    columns: int
    output_width: int
    output_height: int

class HEICFile:
    def __init__(self, filepath: str):
        self.filepath = filepath
        pillow_heif.register_heif_opener()
        
        self._iloc_box: ItemLocationBox | None = None
        self._ftyp_box: Box | None = None
        self._iinf_box: ItemInfoBox | None = None
        self._pitm_box: PrimaryItemBox | None = None
        self._iprp_box: ItemPropertiesBox | None = None
        self._iref_box: ItemReferenceBox | None = None
        self.handler: VendorHandler | None = None
        self.boxes: list[Box] = [] # 顶层盒子

        try:
            with open(self.filepath, 'rb') as f:
                file_size = os.fstat(f.fileno()).st_size
                self.boxes = parse_boxes(f, file_size)
                self._find_essential_boxes(self.boxes)
                self._detect_vendor()
        except Exception as e:
            print(f"CRITICAL ERROR during file parsing: {e}")
            raise

    def _find_essential_boxes(self, boxes: list[Box]):
        """在解析树中查找关键盒子的快捷方式"""
        for box in boxes:
            if isinstance(box, ItemLocationBox): self._iloc_box = box
            if isinstance(box, ItemInfoBox): self._iinf_box = box
            if isinstance(box, PrimaryItemBox): self._pitm_box = box
            if isinstance(box, ItemPropertiesBox): self._iprp_box = box
            if isinstance(box, ItemReferenceBox): self._iref_box = box
            if box.type == 'ftyp': self._ftyp_box = box
            if box.children: self._find_essential_boxes(box.children)

    # --- 辅助方法 (用于修改和构建) ---

    def find_box(self, box_type: str, root_box_list: list[Box] | None = None) -> Box | None:
        """递归查找第一个匹配类型的盒子"""
        if root_box_list is None:
            root_box_list = self.boxes
            
        for box in root_box_list:
            if box.type == box_type:
                return box
            if box.children:
                found = self.find_box(box_type, root_box_list=box.children)
                if found:
                    return found
        return None

    def get_mdat_box(self) -> Box | None:
        """获取顶层的 'mdat' 盒"""
        for box in self.boxes:
            if box.type == 'mdat':
                return box
        return None

    def _remove_box_recursive(self, box_type: str, box_list: list[Box]) -> bool:
        """Helper for remove_box_by_type"""
        for i, box in enumerate(box_list):
            if box.type == box_type:
                box_list.pop(i)
                return True
            if box.children:
                if self._remove_box_recursive(box_type, box.children):
                    return True
        return False

    def remove_box_by_type(self, box_type: str) -> bool:
        """递归查找并移除第一个匹配类型的盒子"""
        return self._remove_box_recursive(box_type, self.boxes)

    # --- START V16 FIX ---
    def remove_item_by_id(self, item_id_to_remove: int):
        """
        (V16 - 垃圾回收版) 
        彻底从 iinf, iloc, ipma, iref 中删除一个 Item ID，
        并从 ipco 中删除孤立的属性，然后重写所有剩余的 ipma 索引。
        """
        print(f"Attempting to remove Item ID {item_id_to_remove} from all references (V16)...")
        
        # 1. 从 iinf (Item Info) 中删除
        if self._iinf_box:
            self._iinf_box.children = [
                c for c in self._iinf_box.children 
                if not (isinstance(c, ItemInfoEntryBox) and c.item_id == item_id_to_remove)
            ]
            self._iinf_box.entries = [
                e for e in self._iinf_box.entries 
                if e.item_id != item_id_to_remove
            ]
            print(f"  - Removed from 'iinf' box.")

        # 2. 从 iloc (Item Location) 中删除
        if self._iloc_box:
            self._iloc_box.locations = [
                loc for loc in self._iloc_box.locations 
                if loc.item_id != item_id_to_remove
            ]
            print(f"  - Removed from 'iloc' box.")

        # 3. (新) 从 iref (Item Reference) 中删除
        if self._iref_box:
            for i in range(len(self._iref_box.children) - 1, -1, -1):
                ref_box = self._iref_box.children[i]
                from_id_size = 4 if self._iref_box.version == 1 else 2
                if len(ref_box.raw_data) >= 4 + from_id_size:
                    from_id = _read_int(ref_box.raw_data, 4, from_id_size)
                    if from_id == item_id_to_remove:
                        self._iref_box.children.pop(i)
                        print(f"  - Removed 'iref' child box (type '{ref_box.type}') with from_id {from_id}.")
            
            ref_types_to_clean = list(self._iref_box.references.keys())
            for ref_type in ref_types_to_clean:
                if item_id_to_remove in self._iref_box.references[ref_type]:
                    del self._iref_box.references[ref_type][item_id_to_remove]
                    print(f"  - Removed from_id {item_id_to_remove} from 'iref.references[{ref_type}]'.")
                
                from_ids_to_clean = list(self._iref_box.references[ref_type].keys())
                for from_id in from_ids_to_clean:
                    self._iref_box.references[ref_type][from_id] = [
                        to_id for to_id in self._iref_box.references[ref_type][from_id]
                        if to_id != item_id_to_remove
                    ]
            print(f"  - Cleaned 'iref.references' of to_id {item_id_to_remove}.")

        # 4. 从 ipma 和 ipco (属性) 中删除
        if self._iprp_box and self._iprp_box.ipma and self._iprp_box.ipco:
            ipma = self._iprp_box.ipma
            ipco = self._iprp_box.ipco
            
            if item_id_to_remove not in ipma.entries:
                print(f"  - Item {item_id_to_remove} not in 'ipma'. No properties to clean.")
                return 
            
            # 4a. 找到要删除的属性索引 (1-based)
            props_to_remove = set(ipma.entries[item_id_to_remove].associations)
            
            # 4b. 从 ipma 中删除该 Item
            del ipma.entries[item_id_to_remove]
            print(f"  - Removed Item {item_id_to_remove} from 'ipma'.")

            # 4c. 确定哪些属性是“孤儿”
            # (即，它们在 props_to_remove 中，但*不在*任何*剩余*的 item 关联中)
            all_remaining_props = set()
            for entry in ipma.entries.values():
                all_remaining_props.update(entry.associations)
            
            orphaned_props = props_to_remove - all_remaining_props
            
            if not orphaned_props:
                print(f"  - No orphaned properties found in 'ipco' to remove.")
                return

            print(f"  - Found orphaned properties to remove from 'ipco': {orphaned_props}")
            
            # 4d. 创建一个“重映射表”
            # 我们从后往前遍历，以安全地删除
            # (索引是 1-based, 但列表是 0-based)
            orphaned_indices_0based = sorted([p - 1 for p in orphaned_props], reverse=True)
            
            # 原始 1-based 索引 -> 新 1-based 索引
            remap_table = {} 
            original_prop_count = len(ipco.children)
            
            # 4e. 从 'ipco.children' 列表中删除孤儿
            for index_0 in orphaned_indices_0based:
                if 0 <= index_0 < len(ipco.children):
                    removed_prop = ipco.children.pop(index_0)
                    print(f"    - Removed property at index {index_0+1} ({removed_prop.type}) from 'ipco'.")
                else:
                    print(f"    - Warning: Orphaned index {index_0+1} out of bounds for 'ipco'.")

            # 4f. 构建重映射表
            # (只有在 ipco 实际发生变化时才需要)
            new_prop_count = len(ipco.children)
            if new_prop_count != original_prop_count:
                print("  - Re-indexing 'ipma' associations...")
                current_new_index = 1
                current_old_index = 1
                orphaned_props_1based = set(i + 1 for i in orphaned_indices_0based)
                
                while current_old_index <= original_prop_count:
                    if current_old_index not in orphaned_props_1based:
                        remap_table[current_old_index] = current_new_index
                        current_new_index += 1
                    current_old_index += 1
                
                # 4g. 应用重映射
                for item_id, entry in ipma.entries.items():
                    new_associations = []
                    for old_prop_index in entry.associations:
                        if old_prop_index in remap_table:
                            new_associations.append(remap_table[old_prop_index])
                        # else:
                        #   该属性已被删除 (不应发生，因为我们只删除了孤儿)
                    
                    # print(f"    - Item {item_id} associations: {entry.associations} -> {new_associations}")
                    entry.associations = new_associations
                print("  - 'ipma' re-indexing complete.")
            else:
                 print("  - No 'ipma' re-indexing needed.")
    # --- END V16 FIX ---

    def set_content_identifier(self, new_content_id: str) -> bool:
        """
        (FIXED) 在 'iinf' 盒中查找主图像的 'infe' 盒，并设置其 item_name
        """
        primary_id = self.get_primary_item_id()
        if not primary_id:
            print("Error: Cannot find primary item ID.")
            return False
            
        if not self._iinf_box:
            print("Error: Cannot find 'iinf' box shortcut (_iinf_box).")
            return False
            
        print(f"Searching for 'infe' box with primary_id = {primary_id}")
        
        # --- START FIX ---
        # 侦测我们日志中看到的 "shifted ID" 格式
        
        target_id = primary_id
        found_ids = [box.item_id for box in self._iinf_box.children if isinstance(box, ItemInfoEntryBox)]
        
        if primary_id not in found_ids:
            print(f"Warning: Primary ID {primary_id} not found directly in 'infe' list.")
            shifted_id = primary_id << 16
            if shifted_id in found_ids:
                print(f"Info: Found vendor-specific shifted ID: {shifted_id} (for {primary_id})")
                target_id = shifted_id
            else:
                print(f"Error: Could not find primary ID {primary_id} OR shifted ID {shifted_id}.")
                print(f"Available 'infe' item IDs found were: {found_ids}")
                return False
        # --- END FIX ---
            
        # 'iinf' 盒的子盒是 'infe' 盒
        for box in self._iinf_box.children:
            if isinstance(box, ItemInfoEntryBox):
                if box.item_id == target_id:
                    print(f"Success: Found 'infe' box for target ID {target_id}. Setting item_name...")
                    box.item_name = new_content_id
                    
                    # 同时更新 self._iinf_box.entries 中的数据以保持一致
                    for entry in self._iinf_box.entries:
                        if entry.item_id == target_id:
                            entry.name = new_content_id
                            break
                    return True
                    
        print(f"Error: Logic failed to find 'infe' box for target ID {target_id} even after check.")
        return False

    # --- 现有的只读方法 (无需修改) ---

    def reconstruct_primary_image(self) -> Image.Image | None:
        """
        使用 pillow-heif 重建主图像，它能自动处理网格图像。
        """
        try:
            print("Reconstructing primary image using pillow-heif...")
            image = Image.open(self.filepath)
            image.load() 
            print("Successfully reconstructed image.")
            return image
        except Exception as e:
            print(f"Failed to reconstruct image with pillow-heif: {e}")
            return None

    def get_primary_image_grid(self) -> Grid | None:
        primary_id = self.get_primary_item_id()
        if not primary_id: return None
        full_size = self.get_image_size(primary_id)
        if not full_size: return None
        full_width, full_height = full_size
        grid_tiles = self.get_grid_layout()
        if not grid_tiles: return None
        first_tile_id = grid_tiles[0]
        tile_size = self.get_image_size(first_tile_id)
        if not tile_size: return None
        tile_width, tile_height = tile_size
        if tile_width == 0 or tile_height == 0: return None
        columns = math.ceil(full_width / tile_width)
        rows = math.ceil(full_height / tile_height)
        return Grid(rows=rows, columns=columns, output_width=full_width, output_height=full_height)

    def get_grid_layout(self) -> list[int] | None:
        primary_id = self.get_primary_item_id()
        if not primary_id: return None
        if not self._iref_box: return None
        
        # 尝试直接匹配
        if 'dimg' in self._iref_box.references and primary_id in self._iref_box.references['dimg']:
            return self._iref_box.references['dimg'].get(primary_id)
        
        # 尝试匹配 "shifted" ID
        shifted_id = primary_id << 16
        if 'dimg' in self._iref_box.references and shifted_id in self._iref_box.references['dimg']:
             print("Info: Using shifted primary ID to find grid layout.")
             return self._iref_box.references['dimg'].get(shifted_id)
             
        return None

    def get_image_size(self, item_id: int) -> tuple[int, int] | None:
        if not (self._iprp_box and self._iprp_box.ipma and self._iprp_box.ipco): return None
        
        target_id = item_id
        if item_id not in self._iprp_box.ipma.entries:
            # 尝试 "shifted" ID
            shifted_id = item_id << 16
            if shifted_id in self._iprp_box.ipma.entries:
                target_id = shifted_id
            else:
                 # 尝试 "unshifted" ID (以防 item_id 本身就是 shifted 的)
                 unshifted_id = item_id & 0x0000FFFF
                 if unshifted_id in self._iprp_box.ipma.entries:
                     target_id = unshifted_id
                 else:
                    # print(f"Warning: Cannot find size associations for item ID {item_id} or variants.")
                    return None
        
        item_associations = self._iprp_box.ipma.entries[target_id].associations
        for assoc in item_associations:
            property_index = assoc - 1
            if 0 <= property_index < len(self._iprp_box.ipco.children):
                prop = self._iprp_box.ipco.children[property_index]
                if isinstance(prop, ImageSpatialExtentsBox):
                    return (prop.image_width, prop.image_height)
        return None

    def get_primary_item_id(self) -> int | None:
        if not self._pitm_box: 
            print("Warning: 'pitm' box not found.")
            return None
        return self._pitm_box.item_id

    def list_items(self):
        if not self._iinf_box: return
        print("Available items in HEIC file:")
        for entry in self._iinf_box.entries: print(f"  - {entry}")

    def _detect_vendor(self):
        if self._ftyp_box:
            compatible_brands = self._ftyp_box.raw_data[4:].decode('ascii', errors='ignore')
            if 'samsung' in compatible_brands.lower(): self.handler = SamsungHandler()
            elif 'apple' in compatible_brands.lower() or 'MiHB' in compatible_brands: self.handler = AppleHandler()
            else: self.handler = VendorHandler()
        else: self.handler = VendorHandler()

    def get_item_data(self, item_id: int) -> bytes | None:
        if not self._iloc_box:
            print("Error: 'iloc' box not found.")
            return None
            
        target_id = item_id
        location = next((loc for loc in self._iloc_box.locations if loc.item_id == target_id), None)
        
        if not location:
            # 尝试检查 "shifted" ID
            shifted_id = item_id << 16
            location = next((loc for loc in self._iloc_box.locations if loc.item_id == shifted_id), None)
            if location:
                print(f"Info: Located item {item_id} using shifted ID {shifted_id} in 'iloc'.")
                target_id = shifted_id
            
        if not location and (item_id & 0xFFFF0000):
            # 尝试反向检查 (如果传入的 ID 已经是 shifted 的)
            unshifted_id = item_id & 0x0000FFFF
            location = next((loc for loc in self._iloc_box.locations if loc.item_id == unshifted_id), None)
            if location:
                 print(f"Info: Located item {item_id} using un-shifted ID {unshifted_id} in 'iloc'.")
                 target_id = unshifted_id

        if not location:
            print(f"Error: Item with ID {item_id} (or variants) not found in 'iloc' box.")
            return None

        if not location.extents:
             print(f"Warning: Item with ID {target_id} has no extents.")
             return b''

        data_chunks = []
        with open(self.filepath, 'rb') as f:
            for offset, length in location.extents:
                f.seek(offset)
                data_chunks.append(f.read(length))
        return b''.join(data_chunks)

    def get_motion_photo_data(self) -> bytes | None:
        if not self.handler: self._detect_vendor()
        offset = self.handler.find_motion_photo_offset(self)
        if offset is not None:
            print("Found motion photo data at offset.")
            with open(self.filepath, 'rb') as f:
                f.seek(offset)
                return f.read()
        return None

    def get_thumbnail_data(self) -> bytes | None:
        print("Attempting to extract thumbnail data...")
        primary_id = self.get_primary_item_id()
        if not primary_id:
            print("Error: Could not determine primary item ID.")
            return None
        
        if not self._iref_box:
            print("Info: No 'iref' box found, cannot search for thumbnail.")
            return None

        if 'thmb' not in self._iref_box.references:
            print("Info: No 'thmb' references found in 'iref' box.")
            return None
        
        target_id = primary_id
        if primary_id not in self._iref_box.references['thmb']:
            # Lпробуйте проверить "shifted" ID
            shifted_primary_id = primary_id << 16
            if shifted_primary_id not in self._iref_box.references['thmb']:
                print(f"Info: Primary item ID {primary_id} (or shifted) has no 'thmb' reference.")
                return None
            
            print("Info: Using shifted primary ID to find thumbnail.")
            target_id = shifted_primary_id


        thumbnail_ids = self._iref_box.references['thmb'][target_id]
        if not thumbnail_ids:
            print(f"Info: Primary item ID {target_id} has 'thmb' reference, but no target IDs.")
            return None

        thumbnail_id = thumbnail_ids[0]
        print(f"Found thumbnail reference: Primary ID {target_id} -> Thumbnail ID {thumbnail_id}")
        
        thumbnail_data = self.get_item_data(thumbnail_id)
        
        if thumbnail_data:
            print(f"Successfully extracted thumbnail data (Item ID {thumbnail_id}).")
            return thumbnail_data
        else:
            print(f"Error: Failed to get data for thumbnail item ID {thumbnail_id}.")
            return None