#!/usr/bin/env python3

"""
Baloo Tools Library
Helper functions to interact directly with the Baloo LMDB index.
"""

import json
import lmdb
import os
import re
import sys
from typing import Tuple

PROPERTIES_ID_MAP = {
    '0': 'Empty',
    '1': 'BitRate',
    '2': 'Channels',
    '3': 'Duration',
    '4': 'Genre',
    '5': 'SampleRate',
    '6': 'TrackNumber',
    '7': 'ReleaseYear',
    '8': 'Comment',
    '9': 'Artist',
    '10': 'Album',
    '11': 'AlbumArtist',
    '12': 'Composer',
    '13': 'Lyricist',
    '14': 'Author',
    '15': 'Title',
    '16': 'Subject',
    '17': 'Generator',
    '18': 'PageCount',
    '19': 'WordCount',
    '20': 'LineCount',
    '21': 'Language',
    '22': 'Copyright',
    '23': 'Publisher',
    '24': 'CreationDate',
    '25': 'Keywords',
    '26': 'Width',
    '27': 'Height',
    '28': 'AspectRatio',
    '29': 'FrameRate',
    '30': 'Manufacturer',
    '31': 'Model',
    '32': 'ImageDateTime',
    '33': 'ImageOrientation',
    '34': 'PhotoFlash',
    '35': 'PhotoPixelXDimension',
    '36': 'PhotoPixelYDimension',
    '37': 'PhotoDateTimeOriginal',
    '38': 'PhotoFocalLength',
    '39': 'PhotoFocalLengthIn35mmFilm',
    '40': 'PhotoExposureTime',
    '41': 'PhotoFNumber',
    '42': 'PhotoApertureValue',
    '43': 'PhotoExposureBiasValue',
    '44': 'PhotoWhiteBalance',
    '45': 'PhotoMeteringMode',
    '46': 'PhotoISOSpeedRatings',
    '47': 'PhotoSaturation',
    '48': 'PhotoSharpness',
    '49': 'PhotoGpsLatitude',
    '50': 'PhotoGpsLongitude',
    '51': 'PhotoGpsAltitude',
    '52': 'TranslationUnitsTotal',
    '53': 'TranslationUnitsWithTranslation',
    '54': 'TranslationUnitsWithDraftTranslation',
    '55': 'TranslationLastAuthor',
    '56': 'TranslationLastUpDate',
    '57': 'TranslationTemplateDate',
    '58': 'OriginUrl',
    '59': 'OriginEmailSubject',
    '60': 'OriginEmailSender',
    '61': 'OriginEmailMessageId',
    '62': 'DiscNumber',
    '63': 'Location',
    '64': 'Performer',
    '65': 'Ensemble',
    '66': 'Arranger',
    '67': 'Conductor',
    '68': 'Opus',
    '69': 'Label',
    '70': 'Compilation',
    '71': 'License',
    '72': 'Rating',
    '73': 'Lyrics',
    '74': 'ReplayGainAlbumPeak',
    '75': 'ReplayGainAlbumGain',
    '76': 'ReplayGainTrackPeak',
    '77': 'ReplayGainTrackGain',
    '78': 'Description',
    '79': 'VideoCodec',
    '80': 'AudioCodec',
    '81': 'PixelFormat',
    '82': 'ColorSpace',
    '83': 'AssistiveAlternateDescription'
}


class BalooTools:
    """Class to interact directly with the Baloo LMDB index."""

    def __init__(self) -> None:
        """Initializes the connection path to the Baloo index."""
        self.baloo_db_path = os.path.join(
            os.path.expanduser("~"), ".local/share/baloo/index"
        )

    def get_info(self, file_id: int) -> json:
        """
        Retrieves file metadata from the Baloo index.

        Args:
            file_id: The integer ID of the file.

        Returns:
            A json with all file metadata fields.
        """
        try:
            # Using context manager ensures the environment is closed properly
            with lmdb.Environment(
                self.baloo_db_path,
                subdir=False,
                readonly=True,
                lock=False,
                max_dbs=20
            ) as env:
                document_data_db = env.open_db(b'documentdatadb')

                with env.begin() as txn:
                    cursor = txn.cursor(document_data_db)

                    # Convert ID to 8-byte little-endian format
                    file_id_bytes = int.to_bytes(
                        file_id, length=8, byteorder='little', signed=False
                    )

                    if cursor.set_range(file_id_bytes):
                        for key, value in cursor:
                            if key != file_id_bytes:
                                break

                            try:
                                jvalue = json.loads(value.decode())
                                return {PROPERTIES_ID_MAP.get(k, k):
                                        v for k, v in jvalue.items()}
                            except (json.JSONDecodeError, KeyError):
                                return {}

        except lmdb.Error as e:
            print(f"Warning: Failed to access Baloo LMDB index: {e}", file=sys.stderr)

        return {}

    def get_resolution(self, file_id: int, sep: str = 'x') -> Tuple[int, int]:
        """
        Retrieves the width and height of an image/video from the Baloo index.

        Args:
            file_id: The integer ID of the file.
            sep: Separator used (unused currently, kept for compatibility).

        Returns:
            A tuple of (width, height) integers. Returns (-1, -1) if not found.
        """
        file_info = self.get_info(file_id)
        try:
            return file_info.get('26', -1), file_info.get('27', -1)
        except (json.JSONDecodeError, KeyError):
            return -1, -1

    def get_tags(self, file_id: int) -> json:
        """
        Retrieves a string with all file tags from the Baloo index.

        Args:
            file_id: The integer ID of the file.

        Returns:
            A json with a field called tags with all tags comma separated.
        """
        try:
            # Using context manager ensures the environment is closed properly
            with lmdb.Environment(
                self.baloo_db_path,
                subdir=False,
                readonly=True,
                lock=False,
                max_dbs=20
            ) as env:
                document_data_db = env.open_db(b'docxatrrterms')

                with env.begin() as txn:
                    cursor = txn.cursor(document_data_db)

                    # Convert ID to 8-byte little-endian format
                    file_id_bytes = int.to_bytes(
                        file_id, length=8, byteorder='little', signed=False
                    )

                    if cursor.set_range(file_id_bytes):
                        for key, value in cursor:
                            if key != file_id_bytes:
                                break

                            text = value.decode('utf-8', errors='replace')
                            text = re.sub(r'\x00(?![T])', '', text)
                            parts = re.split(r'[\x00\x01]', text)

                            tags = []
                            for p in parts:
                                p = p.strip()
                                if p:
                                    tag = p.removeprefix('TAG-').removeprefix('TA')
                                    tags.append(tag)

                            return {'tags': tags}
                            # return {'tags': ",".join(tags)}

        except lmdb.Error as e:
            print(f"Warning: Failed to access Baloo LMDB index: {e}", file=sys.stderr)

        return {}


# Helper function to maintain compatibility with bagheera_search_lib.py
# since it imports `get_resolution` directly.
def get_resolution(file_id: int, sep: str = 'x') -> Tuple[int, int]:
    """Standalone helper function to instantiate BalooTools and get resolution."""
    tools = BalooTools()
    return tools.get_resolution(file_id, sep)


if __name__ == '__main__':
    # CLI execution support for testing
    if len(sys.argv) > 1:
        try:
            target_id = int(sys.argv[1], 16)
            width, height = get_resolution(target_id)
            print(f"{width} {height}")
        except ValueError:
            print("Error: Please provide a valid hexadecimal file ID.", file=sys.stderr)
            sys.exit(1)