"""Cheap FITS header reads and file-integrity checks.

The archive's science keywords (``DEGRADED``, ``ECLIPSE``, ``IMG_MEAN``, ...) live in
the tile-compressed image extension's header, in the first ~18 KB of a 1.7 MB file.
Reading just those bytes and parsing the 80-column cards by hand is ~5x faster than
letting astropy open the file (4.1 ms vs 22.2 ms per file measured on this archive),
which is the difference between a two-hour and a ten-hour pass over the full set.
astropy is still used for pixel data, where it earns its keep.

Every function reports trouble as a structured error string rather than raising, so a
truncated or corrupt file is a *finding* the detector can score, not an exception that
kills a worker.
"""

import logging
import os
import warnings

import numpy as np
from astropy.io import fits
from astropy.utils.exceptions import AstropyUserWarning

from . import db


def quiet_astropy():
    """Silence astropy's per-file chatter about damaged files.

    astropy reports conditions like "File may have been truncated" through both its
    own logger and ``warnings``, either of which floods stderr when scanning
    millions of files -- and damaged files are exactly what a bench run is full of.
    Those conditions are already surfaced per-frame by ``read_header``,
    ``read_image`` and ``verify_datasums``, so nothing is lost by muting them.

    Call once per process; worker pools must call it in each child.
    """
    logging.getLogger("astropy").setLevel(logging.ERROR)
    warnings.simplefilter("ignore", AstropyUserWarning)

#: FITS records are a whole number of 2880-byte blocks.
BLOCK = 2880
CARD = 80
#: How far into a file to look for the image HDU header.  The real ones end by
#: ~17 KB; the cap stops a malformed file from pulling megabytes into memory.
MAX_HEADER_BYTES = BLOCK * 16

#: Sentinel some SUVI products use for undefined pixels.
ZBLANK = -999999


def _parse_value(raw):
    """Parse a FITS card value, honouring quoted strings and comment separators."""
    raw = raw.strip()
    if not raw:
        return None
    if raw[0] == "'":
        # String value: scan for the closing quote, where '' is a literal quote.
        out = []
        i = 1
        while i < len(raw):
            if raw[i] == "'":
                if i + 1 < len(raw) and raw[i + 1] == "'":
                    out.append("'")
                    i += 2
                    continue
                break
            out.append(raw[i])
            i += 1
        return "".join(out).strip()

    # Non-string: everything before the comment separator.
    value = raw.split("/", 1)[0].strip()
    if value in ("T", "F"):
        return value == "T"
    if not value:
        return None
    try:
        return int(value)
    except ValueError:
        pass
    try:
        # FITS permits Fortran-style 'D' exponents.
        return float(value.replace("D", "E").replace("d", "e"))
    except ValueError:
        return value


def _parse_header_blocks(buf, pos):
    """Parse cards from `pos` until the END card.

    Returns (cards, next_pos), or (None, reason) if END is not reached.
    """
    cards = {}
    while pos + BLOCK <= len(buf):
        block = buf[pos : pos + BLOCK]
        pos += BLOCK
        for i in range(0, BLOCK, CARD):
            card = block[i : i + CARD]
            key = card[:8].rstrip()
            if key == b"END":
                return cards, pos
            # Only fixed-format valued cards; COMMENT/HISTORY/CONTINUE lack the '= '.
            if card[8:10] == b"= ":
                try:
                    name = key.decode("ascii").strip()
                    value = _parse_value(card[10:].decode("ascii"))
                except UnicodeDecodeError:
                    continue
                if name and name not in cards:
                    cards[name] = value
    return None, "header has no END card within the scanned region"


def _data_bytes(cards):
    """Size of an HDU's data unit, in bytes, rounded up to whole blocks."""
    naxis = cards.get("NAXIS", 0)
    if not isinstance(naxis, int) or naxis <= 0:
        return 0
    count = 1
    for axis in range(1, naxis + 1):
        length = cards.get(f"NAXIS{axis}")
        if not isinstance(length, int) or length < 0:
            return 0
        count *= length
    bitpix = cards.get("BITPIX", 8)
    if not isinstance(bitpix, int):
        return 0
    size = count * abs(bitpix) // 8
    size += cards.get("PCOUNT", 0) or 0
    return ((size + BLOCK - 1) // BLOCK) * BLOCK


def _is_image_hdu(cards):
    """Whether this HDU carries the image we care about.

    SUVI L2 stores the image tile-compressed inside a BINTABLE, flagged by ZIMAGE;
    its real dimensions are in ZNAXIS, not NAXIS.  A plain IMAGE extension or a
    primary HDU with data also counts.
    """
    if cards.get("ZIMAGE") is True:
        return True
    return isinstance(cards.get("NAXIS"), int) and cards["NAXIS"] > 0


def read_header(path, max_bytes=MAX_HEADER_BYTES):
    """Read the image HDU's header without decompressing pixels.

    Returns (cards, error).  On failure cards is whatever was parsed (possibly
    empty) and error is a short human-readable reason.
    """
    try:
        size = os.path.getsize(path)
        with open(path, "rb") as handle:
            buf = handle.read(max_bytes)
    except OSError as exc:
        return {}, f"unreadable: {exc.strerror or exc}"

    if size == 0:
        return {}, "empty file"
    if len(buf) < BLOCK:
        return {}, f"truncated: {size} bytes is less than one {BLOCK}-byte block"
    if not buf.startswith(b"SIMPLE  ="):
        return {}, "not a FITS file: missing SIMPLE keyword"
    if size % BLOCK:
        # Worth reporting even though the header may still parse: it means the file
        # was cut short mid-transfer, which is exactly the download failure mode.
        truncation = f"truncated: {size} bytes is not a multiple of {BLOCK}"
    else:
        truncation = None

    pos = 0
    while True:
        cards, result = _parse_header_blocks(buf, pos)
        if cards is None:
            return {}, result
        pos = result
        if _is_image_hdu(cards):
            return cards, truncation
        skip = _data_bytes(cards)
        if pos + skip + BLOCK > len(buf):
            return cards, truncation or "no image HDU within the scanned region"
        pos += skip


def header_metadata(cards):
    """Project parsed cards onto the ``header`` table's columns.

    Booleans become 0/1 so they sort and filter naturally in SQL; anything of an
    unexpected type is dropped rather than stored, so a corrupt card cannot poison
    a numeric column.
    """
    values = {}
    for field, keyword, sqltype in db.HEADER_COLUMNS:
        raw = cards.get(keyword)
        if raw is None:
            continue
        if isinstance(raw, bool):
            values[field] = int(raw)
        elif sqltype == "TEXT":
            values[field] = str(raw)
        elif sqltype == "INTEGER":
            values[field] = int(raw) if isinstance(raw, (int, float)) else None
        elif isinstance(raw, (int, float)):
            values[field] = float(raw)
    return {k: v for k, v in values.items() if v is not None}


def scan_header(path):
    """Read and project a file's header in one step.  Returns (values, error)."""
    cards, error = read_header(path)
    return header_metadata(cards), error


def read_image(path):
    """Read the decompressed image array.  Returns (data, error).

    Returns a float32 array with the ZBLANK sentinel converted to NaN, so callers
    do not have to remember that -999999 means "no data" rather than a radiance.

    Detects damage only insofar as astropy refuses to decode it; a file whose image
    HDU survives intact reads fine even if later bytes are missing.  Pair this with
    ``verify_datasums`` when integrity matters rather than availability.
    """
    try:
        with fits.open(path, memmap=False) as hdus:
            data = None
            for hdu in hdus:
                if getattr(hdu, "data", None) is not None and hdu.data.ndim == 2:
                    data = hdu.data
                    break
            if data is None:
                # The blank/corrupt-HDU case filter_FITS.py caught as IndexError.
                return None, "no 2-D image array in file"
    except Exception as exc:  # astropy raises a wide variety on malformed input
        return None, f"{type(exc).__name__}: {exc}"

    data = np.asarray(data, dtype=np.float32)
    blank = data == ZBLANK
    if blank.any():
        data = data.copy()
        data[blank] = np.nan
    return data, None


def datasum(payload):
    """FITS-standard 32-bit ones-complement checksum of a data unit."""
    if len(payload) % 4:
        payload = payload + b"\x00" * (4 - len(payload) % 4)
    total = int(np.frombuffer(payload, dtype=">u4").astype(np.uint64).sum())
    while total >> 32:  # fold the carries back in
        total = (total & 0xFFFFFFFF) + (total >> 32)
    return total


def verify_datasums(path):
    """Check every HDU's stored DATASUM against the bytes on disk.

    Catches silent corruption that a size check misses.  Computed over the raw file,
    deliberately not via astropy: astropy re-serialises a tile-compressed HDU before
    summing, so it reports a mismatch for every healthy SUVI file.

    Returns an error string, or None if all present DATASUMs agree.
    """
    try:
        with open(path, "rb") as handle:
            buf = handle.read()
    except OSError as exc:
        return f"unreadable: {exc.strerror or exc}"

    if len(buf) < BLOCK or not buf.startswith(b"SIMPLE  ="):
        return "not a FITS file"

    pos = 0
    checked = 0
    while pos < len(buf):
        cards, result = _parse_header_blocks(buf, pos)
        if cards is None:
            return result
        size = _data_bytes(cards)
        if result + size > len(buf):
            return "truncated: HDU data unit runs past end of file"
        stored = cards.get("DATASUM")
        if stored is not None:
            try:
                expected = int(str(stored).strip())
            except ValueError:
                return f"malformed DATASUM keyword: {stored!r}"
            if datasum(buf[result : result + size]) != expected:
                name = cards.get("EXTNAME") or ("PRIMARY" if pos == 0 else "extension")
                return f"datasum mismatch in {name}"
            checked += 1
        pos = result + size

    return None if checked else "no DATASUM keywords present"