"""
:mod:`etlplus.file.gz` module.

Helpers for reading/writing GZ files.
"""

from __future__ import annotations

import gzip
from pathlib import Path
from typing import cast

from ..utils.types import JSONData
from ..utils.types import StrPath
from ._archive import infer_archive_payload_format
from ._core_dispatch import read_payload_with_core
from ._core_dispatch import write_payload_with_core
from ._io import ensure_parent_dir
from .base import ArchiveWrapperFileHandlerABC
from .base import ReadOptions
from .base import WriteOptions
from .enums import CompressionFormat
from .enums import FileFormat

# SECTION: EXPORTS ========================================================== #


__all__ = [
    # Classes
    'GzFile',
]


# SECTION: INTERNAL FUNCTIONS =============================================== #


def _resolve_format(
    path: StrPath,
) -> FileFormat:
    """
    Resolve the inner file format from a .gz filename.

    Parameters
    ----------
    path : StrPath
        Path to the GZ file on disk.

    Returns
    -------
    FileFormat
        The inferred inner file format.
    """
    fmt = infer_archive_payload_format(
        path,
        allowed_compressions=(CompressionFormat.GZ,),
        compression_error=f'Not a gzip file: {path}',
    )
    return cast(FileFormat, fmt)


# SECTION: CLASSES ========================================================== #


class GzFile(ArchiveWrapperFileHandlerABC):
    """Handler implementation for GZ files."""

    # -- Class Attributes -- #

    format = FileFormat.GZ
    default_inner_name = 'payload'

    # -- Instance Methods -- #

    def read(
        self,
        path: Path,
        *,
        options: ReadOptions | None = None,
    ) -> JSONData:
        """
        Read GZ content from *path* and parse the inner payload.

        Parameters
        ----------
        path : Path
            Path to the GZ file on disk.
        options : ReadOptions | None, optional
            Optional read parameters.

        Returns
        -------
        JSONData
            Parsed payload.
        """
        fmt = _resolve_format(path)
        payload = self.read_inner_bytes(path, options=options)
        return read_payload_with_core(
            fmt=fmt,
            payload=payload,
            filename=f'payload.{fmt.value}',
        )

    def read_inner_bytes(
        self,
        path: Path,
        *,
        options: ReadOptions | None = None,
    ) -> bytes:
        """
        Read and return decompressed inner payload bytes.

        Parameters
        ----------
        path : Path
            Path to the GZ file on disk.
        options : ReadOptions | None, optional
            Optional read parameters.

        Returns
        -------
        bytes
            Decompressed payload bytes.
        """
        _ = options
        with gzip.open(path, 'rb') as handle:
            return handle.read()

    def write(
        self,
        path: Path,
        data: JSONData,
        *,
        options: WriteOptions | None = None,
    ) -> int:
        """
        Write *data* to GZ at *path* and return record count.

        Parameters
        ----------
        path : Path
            Path to the GZ file on disk.
        data : JSONData
            Data to write.
        options : WriteOptions | None, optional
            Optional write parameters.

        Returns
        -------
        int
            Number of records written.
        """
        fmt = _resolve_format(path)
        count, payload = write_payload_with_core(
            fmt=fmt,
            data=data,
            filename=f'payload.{fmt.value}',
        )

        self.write_inner_bytes(path, payload, options=options)
        return count

    def write_inner_bytes(
        self,
        path: Path,
        payload: bytes,
        *,
        options: WriteOptions | None = None,
    ) -> None:
        """
        Compress and write inner payload bytes.

        Parameters
        ----------
        path : Path
            Path to the GZ file on disk.
        payload : bytes
            Raw inner payload bytes.
        options : WriteOptions | None, optional
            Optional write parameters.
        """
        _ = options
        ensure_parent_dir(path)
        with gzip.open(path, 'wb') as handle:
            handle.write(payload)
