# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
``fitscheck`` is a command line script based on astropy.io.fits for verifying
and updating the CHECKSUM and DATASUM keywords of .fits files.  ``fitscheck``
can also detect and often fix other FITS standards violations.  ``fitscheck``
facilitates re-writing the non-standard checksums originally generated by
astropy.io.fits with standard checksums which will interoperate with CFITSIO.

``fitscheck`` will refuse to write new checksums if the checksum keywords are
missing or their values are bad.  Use ``--force`` to write new checksums
regardless of whether or not they currently exist or pass.  Use
``--ignore-missing`` to tolerate missing checksum keywords without comment.

Example uses of fitscheck:

1. Add checksums::

    $ fitscheck --write *.fits

2. Write new checksums, even if existing checksums are bad or missing::

    $ fitscheck --write --force *.fits

3. Verify standard checksums and FITS compliance without changing the files::

    $ fitscheck --compliance *.fits

4. Only check and fix compliance problems,  ignoring checksums::

    $ fitscheck --checksum none --compliance --write *.fits

5. Verify standard interoperable checksums::

    $ fitscheck *.fits

6. Delete checksum keywords::

    $ fitscheck --checksum remove --write *.fits

"""

import argparse
import logging
import sys
import warnings

from astropy import __version__
from astropy.io import fits

log = logging.getLogger("fitscheck")

DESCRIPTION = """
e.g. fitscheck example.fits

Verifies and optionally re-writes the CHECKSUM and DATASUM keywords
for a .fits file.
Optionally detects and fixes FITS standard compliance problems.

This script is part of the Astropy package. See
https://docs.astropy.org/en/latest/io/fits/usage/scripts.html#module-astropy.io.fits.scripts.fitscheck
for further documentation.
""".strip()


def handle_options(args):
    if not len(args):
        args = ["-h"]

    parser = argparse.ArgumentParser(
        description=DESCRIPTION, formatter_class=argparse.RawDescriptionHelpFormatter
    )

    parser.add_argument(
        "--version", action="version", version=f"%(prog)s {__version__}"
    )

    parser.add_argument(
        "fits_files", metavar="file", nargs="+", help=".fits files to process."
    )

    parser.add_argument(
        "-k",
        "--checksum",
        dest="checksum_kind",
        choices=["standard", "remove", "none"],
        help="Choose FITS checksum mode or none.  Defaults standard.",
        default="standard",
    )

    parser.add_argument(
        "-w",
        "--write",
        dest="write_file",
        help="Write out file checksums and/or FITS compliance fixes.",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "-f",
        "--force",
        dest="force",
        help="Do file update even if original checksum was bad.",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "-c",
        "--compliance",
        dest="compliance",
        help="Do FITS compliance checking; fix if possible.",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "-i",
        "--ignore-missing",
        dest="ignore_missing",
        help="Ignore missing checksums.",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "-v",
        "--verbose",
        dest="verbose",
        help="Generate extra output.",
        default=False,
        action="store_true",
    )

    global OPTIONS
    OPTIONS = parser.parse_args(args)

    if OPTIONS.checksum_kind == "none":
        OPTIONS.checksum_kind = False
    elif OPTIONS.checksum_kind == "standard":
        OPTIONS.checksum_kind = True
    elif OPTIONS.checksum_kind == "remove":
        OPTIONS.write_file = True
        OPTIONS.force = True

    return OPTIONS.fits_files


def setup_logging():
    log.handlers.clear()

    if OPTIONS.verbose:
        log.setLevel(logging.INFO)
    else:
        log.setLevel(logging.WARNING)

    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter("%(message)s"))
    log.addHandler(handler)


def verify_checksums(filename):
    """
    Prints a message if any HDU in `filename` has a bad checksum or datasum.
    """
    with warnings.catch_warnings(record=True) as wlist:
        warnings.simplefilter("always")
        with fits.open(filename, checksum=OPTIONS.checksum_kind) as hdulist:
            for i, hdu in enumerate(hdulist):
                # looping on HDUs is needed to read them and verify the
                # checksums
                if not OPTIONS.ignore_missing:
                    if not hdu._checksum:
                        log.warning(
                            f"MISSING {filename!r} .. Checksum not found in HDU #{i}"
                        )
                        return 1
                    if not hdu._datasum:
                        log.warning(
                            f"MISSING {filename!r} .. Datasum not found in HDU #{i}"
                        )
                        return 1

    for w in wlist:
        if str(w.message).startswith(
            ("Checksum verification failed", "Datasum verification failed")
        ):
            log.warning("BAD %r %s", filename, str(w.message))
            return 1

    log.info(f"OK {filename!r}")
    return 0


def verify_compliance(filename):
    """Check for FITS standard compliance."""
    with fits.open(filename) as hdulist:
        try:
            hdulist.verify("exception")
        except fits.VerifyError as exc:
            log.warning("NONCOMPLIANT %r .. %s", filename, str(exc).replace("\n", " "))
            return 1
    return 0


def update(filename):
    """
    Sets the ``CHECKSUM`` and ``DATASUM`` keywords for each HDU of `filename`.

    Also updates fixes standards violations if possible and requested.
    """
    output_verify = "silentfix" if OPTIONS.compliance else "ignore"

    # For unit tests we reset temporarily the warning filters. Indeed, before
    # updating the checksums, fits.open will verify the existing checksums and
    # raise warnings, which are later caught and converted to log.warning...
    # which is an issue when testing, using the "error" action to convert
    # warnings to exceptions.
    with warnings.catch_warnings():
        warnings.resetwarnings()
        with fits.open(
            filename,
            do_not_scale_image_data=True,
            checksum=OPTIONS.checksum_kind,
            mode="update",
        ) as hdulist:
            hdulist.flush(output_verify=output_verify)


def process_file(filename):
    """
    Handle a single .fits file,  returning the count of checksum and compliance
    errors.
    """
    try:
        checksum_errors = verify_checksums(filename)
        if OPTIONS.compliance:
            compliance_errors = verify_compliance(filename)
        else:
            compliance_errors = 0
        if OPTIONS.write_file and checksum_errors == 0 or OPTIONS.force:
            update(filename)
        return checksum_errors + compliance_errors
    except Exception as e:
        log.error(f"EXCEPTION {filename!r} .. {e}")
        return 1


def main(args=None):
    """
    Processes command line parameters into options and files,  then checks
    or update FITS DATASUM and CHECKSUM keywords for the specified files.
    """
    errors = 0
    fits_files = handle_options(args or sys.argv[1:])
    setup_logging()
    for filename in fits_files:
        errors += process_file(filename)
    if errors:
        log.warning(f"{errors} errors")
    return int(bool(errors))