HelioParser/utils/metadata_parser.py

"""
Парсер для извлечения FITS-метаданных из JP2 файлов
"""

import xml.etree.ElementTree as ET
from typing import Dict, Optional
import struct


class MetadataParser:
    """Парсер FITS-метаданных из JP2 файлов"""

    @staticmethod
    def extract_metadata(filepath: str) -> Optional[Dict[str, str]]:
        """
        Извлекает FITS-метаданные из JP2 файла

        Args:
            filepath: Путь к JP2 файлу

        Returns:
            Словарь с метаданными или None
        """
        try:
            # Пытаемся прочитать JP2 файл как бинарный и найти XML
            with open(filepath, 'rb') as f:
                data = f.read()

            # Ищем XML данные в файле (JP2 может содержать XML в специальных боксах)
            # Простой способ: ищем теги XML
            xml_start = data.find(b'<?xml')
            if xml_start == -1:
                xml_start = data.find(b'<fits')

            if xml_start != -1:
                # Ищем конец XML
                xml_end = data.find(b'>', xml_start + 100)  # Примерный поиск
                # Находим закрывающий тег
                if xml_end != -1:
                    # Расширяем поиск до полного XML
                    depth = 1
                    pos = xml_start
                    while depth > 0 and pos < len(data):
                        pos += 1
                        if pos >= len(data):
                            break
                        if data[pos:pos+2] == b'</':
                            depth -= 1
                        elif data[pos:pos+1] == b'<':
                            depth += 1

                    if pos < len(data):
                        xml_data = data[xml_start:pos+1].decode('utf-8', errors='ignore')
                        return MetadataParser._parse_fits_xml(xml_data)

            # Если не нашли XML, пытаемся извлечь из текстовых блоков
            # Ищем ASCII текст
            text = data.decode('utf-8', errors='ignore')

            # Ищем FITS-подобные ключи
            metadata = {}
            fits_keys = ['TELESCOP', 'INSTRUME', 'WAVELNTH', 'DATE-OBS', 'EXPTIME',
                        'CRPIX1', 'CRPIX2', 'CDELT1', 'CDELT2', 'NAXIS1', 'NAXIS2']

            for key in fits_keys:
                # Ищем ключ в тексте
                pattern = f'{key}='
                start = text.find(pattern)
                if start != -1:
                    # Находим значение
                    value_start = start + len(pattern)
                    # Ищем конец строки или следующий ключ
                    value_end = text.find('/', value_start)
                    if value_end == -1:
                        value_end = text.find('\n', value_start)
                    if value_end == -1:
                        value_end = text.find(' ', value_start + 20)

                    value = text[value_start:value_end].strip().strip("'\"")
                    if value:
                        metadata[key] = value

            if metadata:
                return MetadataParser._format_metadata(metadata)

            return None

        except Exception as e:
            print(f"Ошибка извлечения метаданных: {e}")
            return None

    @staticmethod
    def _parse_fits_xml(xml_content: str) -> Dict[str, str]:
        """
        Парсит XML с FITS-заголовком

        Args:
            xml_content: XML строка

        Returns:
            Словарь с параметрами FITS
        """
        metadata = {}

        try:
            root = ET.fromstring(xml_content)

            # Ищем секцию fits или FITS
            fits_section = root.find('.//fits') or root.find('.//FITS')

            if fits_section is not None:
                for child in fits_section:
                    # Извлекаем ключ и значение
                    key = child.tag
                    value = child.text if child.text else ""

                    # Убираем namespace если есть
                    if '}' in key:
                        key = key.split('}')[-1]

                    metadata[key.upper()] = value.strip()

            # Если не нашли fits секцию, ищем другие возможные места
            if not metadata:
                for child in root.iter():
                    if child.tag.endswith('keyword'):
                        key = child.get('name', '')
                        value = child.text if child.text else ''
                        if key:
                            metadata[key.upper()] = value.strip()

            # Форматируем некоторые ключи для удобства чтения
            metadata = MetadataParser._format_metadata(metadata)

        except Exception as e:
            print(f"Ошибка парсинга XML: {e}")

        return metadata

    @staticmethod
    def _format_metadata(metadata: Dict[str, str]) -> Dict[str, str]:
        """
        Форматирует метаданные для удобного отображения

        Args:
            metadata: Сырые метаданные

        Returns:
            Отформатированные метаданные
        """
        formatted = {}

        # Переименовываем некоторые ключи для понятности
        key_mapping = {
            'TELESCOP': 'Телескоп',
            'INSTRUME': 'Инструмент',
            'DETECTOR': 'Детектор',
            'WAVELNTH': 'Длина волны',
            'WAVEUNIT': 'Единица длины волны',
            'DATE-OBS': 'Дата наблюдения',
            'DATE-BEG': 'Начало экспозиции',
            'DATE-END': 'Конец экспозиции',
            'EXPTIME': 'Время экспозиции (сек)',
            'CRPIX1': 'Центр X (пикс)',
            'CRPIX2': 'Центр Y (пикс)',
            'CDELT1': 'Шаг пикселя X (arcsec)',
            'CDELT2': 'Шаг пикселя Y (arcsec)',
            'CROTA2': 'Угол поворота (град)',
            'NAXIS1': 'Ширина (пикс)',
            'NAXIS2': 'Высота (пикс)',
            'BITPIX': 'Бит на пиксель',
            'BZERO': 'Смещение данных',
            'BSCALE': 'Масштаб данных',
            'IMG_TYPE': 'Тип изображения',
            'QUALITY': 'Качество',
            'LEVEL': 'Уровень обработки',
            'STATUS': 'Статус',
            'OBSRVTRY': 'Обсерватория'
        }

        for key, value in metadata.items():
            # Используем переименованный ключ или оригинал
            display_key = key_mapping.get(key, key)
            formatted[display_key] = value

        return formatted

    @staticmethod
    def extract_from_jp2_box(filepath: str) -> Optional[Dict]:
        """
        Альтернативный метод: извлечение метаданных через чтение JP2 боксов
        Не требует glymur, читает файл напрямую
        """
        try:
            with open(filepath, 'rb') as f:
                data = f.read()

            # JP2 signature
            if data[0:4] != b'\x00\x00\x00\x0c':
                return None

            # Ищем XML бокс (box type 'xml ')
            offset = 0
            while offset < len(data) - 8:
                box_len = struct.unpack('>I', data[offset:offset+4])[0]
                box_type = data[offset+4:offset+8]

                if box_type == b'xml ' or box_type == b'XML ':
                    # Нашли XML бокс
                    xml_data = data[offset+8:offset+box_len]
                    try:
                        xml_str = xml_data.decode('utf-8', errors='ignore')
                        return MetadataParser._parse_fits_xml(xml_str)
                    except:
                        pass

                if box_len == 0:
                    break
                offset += box_len

            return None

        except Exception as e:
            print(f"Ошибка чтения JP2 боксов: {e}")
            return None