from __future__ import annotations

"""Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Identity and gzip compression helpers for the transfer layer.
"""

from gzip import GzipFile
from shutil import copyfileobj

from mcp_proxy_adapter.transfer.errors import TransferCompressionError

ALLOWED_ALGORITHMS: tuple[str, ...] = ("identity", "gzip")


def validate_compression_algorithm(algorithm: str) -> str:
    if algorithm not in ALLOWED_ALGORITHMS:
        raise TransferCompressionError(
            message=f"Unsupported compression algorithm: {algorithm!r}",
            compression=algorithm,
            phase="validate",
        )
    return algorithm


def compress_file(source_path: str, destination_path: str, algorithm: str) -> str:
    validate_compression_algorithm(algorithm)
    if algorithm == "identity":
        try:
            with open(source_path, "rb") as src, open(destination_path, "wb") as dst:
                copyfileobj(src, dst)
        except OSError as exc:
            raise TransferCompressionError(
                message=f"identity compression failed: {exc}",
                compression="identity",
                phase="compress",
            ) from exc
    elif algorithm == "gzip":
        try:
            with open(source_path, "rb") as src:
                with open(destination_path, "wb") as dst:
                    with GzipFile(fileobj=dst, mode="wb") as gz:
                        copyfileobj(src, gz)
        except OSError as exc:
            raise TransferCompressionError(
                message=f"gzip compression failed: {exc}",
                compression="gzip",
                phase="compress",
            ) from exc
    else:
        raise TransferCompressionError(
            message=f"Unsupported compression algorithm: {algorithm!r}",
            compression=algorithm,
            phase="validate",
        )
    return destination_path


def decompress_file(source_path: str, destination_path: str, algorithm: str) -> str:
    validate_compression_algorithm(algorithm)
    if algorithm == "identity":
        try:
            with open(source_path, "rb") as src, open(destination_path, "wb") as dst:
                copyfileobj(src, dst)
        except OSError as exc:
            raise TransferCompressionError(
                message=f"identity decompression failed: {exc}",
                compression="identity",
                phase="decompress",
            ) from exc
    elif algorithm == "gzip":
        try:
            with open(source_path, "rb") as src:
                with GzipFile(fileobj=src, mode="rb") as gz:
                    with open(destination_path, "wb") as dst:
                        copyfileobj(gz, dst)
        except OSError as exc:
            raise TransferCompressionError(
                message=f"gzip decompression failed: {exc}",
                compression="gzip",
                phase="decompress",
            ) from exc
    else:
        raise TransferCompressionError(
            message=f"Unsupported compression algorithm: {algorithm!r}",
            compression=algorithm,
            phase="validate",
        )
    return destination_path
