"""S3 file downloader with integrity verification and validation.

Downloads files from S3 with comprehensive pre- and post-download validation including
size limits, disk space checks, and MD5 checksum verification for data integrity.

Validation Flow:
    1. Check file exists and size is within limits
    2. Verify sufficient disk space available
    3. Download file from S3
    4. Verify downloaded file size matches S3 metadata
    5. Verify MD5 checksum against S3 ETag (single-part uploads only)

Note: MD5 verification is skipped for multipart uploads as S3 ETags are composite hashes.
"""

from __future__ import annotations

import os
from pathlib import Path

from abacus_common_logic.utils.formatting import format_bytes

from src.connectors.s3 import S3Connection
from src.errors import (
    ChecksumMismatchError,
    FileIntegrityError,
    FileSizeExceededError,
    FileSystemError,
    S3FileNotFoundError,
)
from src.infra.log import logger
from src.utils.file_utils import (
    compute_md5,
    get_dir_path,
    is_simple_etag,
    try_disk_space,
)


class S3Downloader:
    """S3 downloader with integrity verification."""

    def __init__(self, s3_connection: S3Connection, max_bytes: int | None):
        """Initialize downloader with S3 connection and optional size limit.

        Args:
            s3_connection: S3 connector instance.
            max_bytes: Maximum allowed file size in bytes, or None for no limit.
        """
        self._s3_connection = s3_connection
        self._max_bytes = max_bytes
        if max_bytes is not None and max_bytes < 0:
            raise ValueError('max_bytes cannot be negative')

    def download(self, s3_bucket: str, s3_key: str, file_path: str | Path) -> int:
        """Download file from S3 with pre- and post-download validation.

        Args:
            s3_bucket: S3 bucket name.
            s3_key: S3 object key.
            file_path: Local destination path.

        Returns:
            File size in bytes.

        Raises:
            S3FileNotFoundError: If file is not found in S3.
            FileSizeExceededError: If file exceeds maximum size.
            FileIntegrityError: If downloaded file size doesn't match expected.
            FileSystemError: If downloaded file cannot be verified.
            ChecksumMismatchError: If MD5 checksum verification fails.
            TransientError: If disk space is insufficient.
        """
        logger.info(f'Checking file exists at: s3://{s3_bucket}/{s3_key}')
        metadata = self._s3_connection.get_file_metadata(s3_bucket, s3_key)
        if not metadata:
            raise S3FileNotFoundError(
                f'S3 file not found: s3://{s3_bucket}/{s3_key}. '
                'File may have been deleted or moved'
            )

        file_size = metadata.size
        logger.info(f'S3 file size: {format_bytes(file_size)}')
        if self._max_bytes and file_size > self._max_bytes:
            raise FileSizeExceededError(
                f'File size exceeds maximum: '
                f'{format_bytes(file_size)} > '
                f'{format_bytes(self._max_bytes)}'
            )

        logger.info('Checking disk space')
        try_disk_space(file_size, get_dir_path(file_path))

        logger.info(f'Downloading file to: {file_path}')
        self._s3_connection.download_file(s3_bucket, s3_key, file_path)
        logger.info('File downloaded')

        logger.info('Verifying download')
        self._verify_file_size(file_path, file_size)
        self._verify_file_checksum(file_path, metadata.etag)
        logger.info('Download verified')

        return file_size

    def _cleanup_file(self, file_path: str | Path) -> None:
        """Remove file if it exists, suppressing errors.

        Args:
            file_path: Path to file to remove.
        """
        try:
            os.remove(file_path)
            logger.info(f'Cleaned up partial/corrupt file: {file_path}')
        except FileNotFoundError:
            logger.debug(f'File already removed: {file_path}')
        except OSError as e:
            logger.error(f'Failed to cleanup file {file_path}: {e}')

    def _verify_file_checksum(self, file_path: str | Path, etag: str) -> None:
        """Verify downloaded file checksum against S3 ETag.

        For single-part uploads, verifies MD5 checksum.
        For multipart uploads, logs warning (ETag is not simple MD5).

        Args:
            file_path: Path to downloaded file.
            etag: S3 ETag value.

        Raises:
            ChecksumMismatchError: If MD5 verification fails.
            FileSystemError: If checksum cannot be computed.
        """
        if not is_simple_etag(etag):
            logger.warning(
                'Skipping checksum verification: file was uploaded as multipart'
            )
            return

        logger.debug('Computing MD5 checksum')
        try:
            file_md5 = compute_md5(file_path)
        except OSError as e:
            logger.error(f'Failed to compute checksum: {e}')
            self._cleanup_file(file_path)
            raise FileSystemError(
                f'Cannot compute checksum for {file_path}: {e}'
            ) from e

        if file_md5 != etag:
            logger.error(f'Checksum mismatch: expected {etag}, got {file_md5}')
            self._cleanup_file(file_path)
            raise ChecksumMismatchError(
                f'MD5 checksum mismatch: expected {etag}, got {file_md5}'
            )

    def _verify_file_size(self, file_path: str | Path, file_size: int) -> None:
        """Verify downloaded file size matches expected size.

        Args:
            file_path: Path to downloaded file.
            file_size: Expected file size in bytes from S3 metadata.

        Raises:
            FileIntegrityError: If file size doesn't match expected size.
            FileSystemError: If file size cannot be determined.
        """
        try:
            local_file_size = os.path.getsize(file_path)
        except OSError as e:
            logger.error(f'Failed to verify download: {e}')
            self._cleanup_file(file_path)
            raise FileSystemError(f'Cannot verify download at {file_path}: {e}') from e

        if file_size != local_file_size:
            logger.error(
                f'Download integrity check failed: '
                f'expected {format_bytes(file_size)}, '
                f'got {format_bytes(local_file_size)}'
            )
            self._cleanup_file(file_path)
            raise FileIntegrityError(
                f'Download size mismatch: '
                f'expected {file_size} bytes, '
                f'got {local_file_size} bytes'
            )