from hashlib import sha256, md5
from pathlib import Path

import aiofiles


class HashCheckError(ValueError):
    def __init__(
        self,
        msg: str,
        hash_type: str,
        expected_hash: str,
        actual_hash: str,
        bytes_read: int,
    ):
        super().__init__(msg)

        self.hash_type = hash_type
        self.expected_hash = expected_hash
        self.actual_hash = actual_hash
        self.bytes_read = bytes_read

    def __str__(self):
        return f"{self.args[0]} (Expected: {self.expected_hash} Actual: {self.actual_hash})"


async def validate_checksum(file: Path, sha256_expected, md5_expected):
    chunk_read_size = 10 * 1024 * 1024  # 10MB reads
    if sha256_expected is not None:
        file_hasher = sha256()
        expected_checksum = sha256_expected
        checksum_type = "SHA256"
    else:
        file_hasher = md5()
        expected_checksum = md5_expected
        checksum_type = "MD5"

    bytes_read = 0
    async with aiofiles.open(file, "rb") as written_file:
        while True:
            file_bytes = await written_file.read(chunk_read_size)
            if not file_bytes:
                break
            file_hasher.update(file_bytes)
            bytes_read += len(file_bytes)

    actual_checksum = file_hasher.hexdigest()

    if actual_checksum != expected_checksum:
        raise HashCheckError(
            "checksum failure",
            checksum_type,
            expected_checksum,
            actual_checksum,
            bytes_read,
        )
