#!/usr/bin/env python3

import argparse
import hashlib
import sys
from pathlib import Path

DEFAULT_BLOCK_SIZE = 1024 * 1024  # 1 MB
SUPPORTED_ALGOS = ["md5", "sha1", "sha224", "sha256"]

def compute_hash(path: Path, algorithm: str, block_size: int = DEFAULT_BLOCK_SIZE) -> str:
    """Compute the hash of a file using the chosen algorithm."""
    try:
        h = hashlib.new(algorithm)
    except ValueError:
        raise SystemExit(f"Error: Unsupported hash algorithm '{algorithm}'")

    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(block_size), b""):
            h.update(chunk)

    return h.hexdigest()

def main():
    parser = argparse.ArgumentParser(
        description="Compute hash(es) of a file or verify against a known hash."
    )
    parser.add_argument("file", type=Path, help="Path to the file to hash")
    parser.add_argument(
        "known_hash",
        nargs="?",
        help="Optional known hash to verify the file against",
    )

    parser.add_argument("-md5", action="store_true", help="Use MD5 algorithm")
    parser.add_argument("-sha1", action="store_true", help="Use SHA-1 algorithm")
    parser.add_argument("-sha224", action="store_true", help="Use SHA-224 algorithm")
    parser.add_argument("-sha256", action="store_true", help="Use SHA-256 algorithm")

    args = parser.parse_args()

    if not args.file.exists() or not args.file.is_file():
        print(f"Error: File not found -> {args.file}", file=sys.stderr)
        sys.exit(1)

    selected_algos = [
        algo
        for algo, used in {
            "md5": args.md5,
            "sha1": args.sha1,
            "sha224": args.sha224,
            "sha256": args.sha256,
        }.items()
        if used
    ]

    if not selected_algos:
        selected_algos = ["sha1"] if args.known_hash else SUPPORTED_ALGOS

    print()
    if args.known_hash:
        # Verification mode
        algo = selected_algos[0]
        digest = compute_hash(args.file, algo)
        print(f"{algo.upper()} hash: {digest}")
        if digest.lower() == args.known_hash.lower():
            print("\033[32mCorrect\033[0m")
            print("")
            sys.exit(0)
        else:
            print("\033[31mIncorrect\033[0m")
            print("")
            sys.exit(1)
    else:
        # Normal hashing mode
        print("Algorithm       Hash")
        print("---------       ----")
        for algo in selected_algos:
            digest = compute_hash(args.file, algo)
            print(f"{algo.upper():<8}        {digest}")
        print()

if __name__ == "__main__":
    main()