import hashlib
import os
import tarfile
import tempfile

from occystrap import constants
from occystrap.filters.base import ImageFilter
from occystrap.tarformat import select_tar_format_for_layer
from shakenfist_utilities import logs


LOG = logs.setup_console(__name__)


class TimestampNormalizer(ImageFilter):
    """Normalizes timestamps in image layers for reproducible builds.

    This filter rewrites layer tarballs to set all file modification times
    to a consistent value (default: 0, Unix epoch). Since this changes the
    layer content, the SHA256 hash is recalculated and the layer name is
    updated to match.

    This is useful for creating reproducible image tarballs where the same
    source content always produces the same output, regardless of when the
    files were originally created or modified.
    """

    def __init__(self, wrapped_output, timestamp=0,
                 temp_dir=None, diff_id_map=None):
        """Initialize the timestamp normalizer.

        Args:
            wrapped_output: The ImageOutput to pass normalized elements
                to.
            timestamp: The Unix timestamp to set for all files
                (default: 0).
            temp_dir: Directory for temporary files (default:
                system temp directory).
            diff_id_map: Shared dict for cross-image diff_id tracking
                (default: None).
        """
        super().__init__(
            wrapped_output, temp_dir=temp_dir,
            diff_id_map=diff_id_map)
        self.timestamp = timestamp

    def _normalize_layer(self, layer_data):
        """Normalize timestamps in a layer tarball.

        Creates a new tarball with all timestamps set to self.timestamp,
        calculates the new SHA256 hash, and returns both.

        Uses USTAR format when possible (smaller output), falls back to
        PAX format when layer contents require it. See tarformat.py.

        Args:
            layer_data: File-like object containing the original layer.

        Returns:
            Tuple of (normalized_file_handle, new_sha256_hex)
        """
        # Determine optimal tar format based on transformed members
        def transform(member):
            member.mtime = self.timestamp
            return member

        tar_format = select_tar_format_for_layer(layer_data, transform)

        with tempfile.NamedTemporaryFile(
                delete=False, dir=self.temp_dir) as normalized_tf:
            try:
                # Create a new tarball with normalized timestamps
                with tarfile.open(fileobj=normalized_tf, mode='w',
                                  format=tar_format) as normalized_tar:
                    layer_data.seek(0)
                    with tarfile.open(fileobj=layer_data, mode='r') as \
                            layer_tar:
                        for member in layer_tar:
                            # Normalize all timestamp fields
                            member.mtime = self.timestamp

                            # Extract the file data if it's a regular file
                            if member.isfile():
                                fileobj = layer_tar.extractfile(member)
                                normalized_tar.addfile(member, fileobj)
                            else:
                                normalized_tar.addfile(member)

                # Calculate SHA256 of the normalized tarball
                normalized_tf.flush()
                normalized_tf.seek(0)
                h = hashlib.sha256()
                while True:
                    chunk = normalized_tf.read(8192)
                    if not chunk:
                        break
                    h.update(chunk)

                new_sha = h.hexdigest()

                # Return a new file handle and the hash
                normalized_tf.seek(0)
                return open(normalized_tf.name, 'rb'), new_sha

            except Exception:
                os.unlink(normalized_tf.name)
                raise

    def process_image_element(self, element):
        """Process an image element, normalizing layer
        timestamps.

        Config files are buffered so diff_ids can be
        updated in finalize(). Layers have their
        timestamps normalized and their names updated
        to reflect the new SHA256 hash.
        """
        if element.element_type == constants.CONFIG_FILE:
            self._buffer_config(element)
        elif (element.element_type
                == constants.IMAGE_LAYER
                and element.data is not None):
            LOG.debug(
                'Normalizing timestamps in layer %s'
                % element.name)
            normalized_data, new_name = \
                self._normalize_layer(element.data)
            self._record_new_diff_id(
                new_name, element.layer_index,
                original_hex=element.name)

            try:
                self._wrapped.process_image_element(
                    constants.ImageElement(
                        element.element_type,
                        new_name,
                        normalized_data,
                        layer_index=(
                            element.layer_index)))
            finally:
                try:
                    normalized_data.close()
                    os.unlink(normalized_data.name)
                except Exception:
                    pass
        else:
            if (element.element_type
                    == constants.IMAGE_LAYER):
                self._skip_layer(element.layer_index)
            self._wrapped.process_image_element(
                element)
