"""Versioned staging and retention for the endpoint agent binaries.

The site-agent packager ships a prebuilt shopdbagent.exe and an optional
shopdb-endpoint-tray.exe. Originally each was a single staged file under
instance/siteagent/; this module adds per-version staging so an operator can
drop in a new build, keep the last few, and roll a push out from a known
version.

Layout (all under the Flask instance directory):

    siteagent/
      shopdbagent.exe              legacy flat path (pre-versioning)
      shopdb-endpoint-tray.exe     legacy flat path (pre-versioning)
      agent/
        manifest.json              {current, versions:[{version,uploadedat,size}]}
        1.6.3/shopdbagent.exe
        1.6.2/shopdbagent.exe
      tray/
        manifest.json
        1.0.0/shopdb-endpoint-tray.exe

Retention: the RETAIN most recently uploaded versions per binary are kept; the
current version is never pruned even if it falls outside that window. The legacy
flat file is treated as an unmanaged fallback: the packager still ships it when
no versioned build exists, so an install that predates versioning keeps working.

This module holds NO routes and NO permission checks - it is pure staging logic
so it can be unit-tested without a request context and reused by both the upload
routes and the packager. Callers own authorization.
"""

import hashlib
import io
import json
import os
import re
import struct
import zipfile
from datetime import datetime

from flask import current_app

# How many versions to retain per binary. The current version is exempt from
# pruning, so the effective floor is max(RETAIN, 1-including-current).
RETAIN = 3

# Per-binary staging spec. 'exe' is the canonical filename shipped in the ZIP and
# stored inside each version directory; 'legacy' is the flat pre-versioning path
# under siteagent/ kept as a fallback source.
BINARY_SPECS = {
    'agent': {'exe': 'shopdbagent.exe', 'legacy': 'shopdbagent.exe'},
    'tray': {'exe': 'shopdb-endpoint-tray.exe', 'legacy': 'shopdb-endpoint-tray.exe'},
}

# A dotted numeric version, 1 to 4 parts (matches .NET assembly versions like
# 1.6.3 or 1.6.3.0). Deliberately strict so a typo or an arbitrary label cannot
# become a directory name.
_VERSION_RE = re.compile(r'^\d+(\.\d+){0,3}$')

# Windows PE magic. Uploaded agent binaries must be real PE executables; a quick
# header check rejects an operator who uploads the wrong file (a ZIP, a script)
# before it is ever staged or pushed.
_PE_MAGIC = b'MZ'


def is_valid_binary_key(binkey):
    """Whether binkey names a managed binary (agent or tray)."""
    return binkey in BINARY_SPECS


def is_valid_version(version):
    """Whether version is a strict dotted-numeric string safe to use as a dir."""
    return bool(version) and bool(_VERSION_RE.match(version))


def looks_like_pe(data):
    """Whether data begins with the Windows PE 'MZ' magic."""
    return len(data) >= 2 and data[:2] == _PE_MAGIC


# RT_VERSION resource type id (the version resource Windows shows in file
# properties). We must read the version from THIS resource on the outer
# executable - not the first VS_FIXEDFILEINFO signature we can find in the bytes.
# A self-contained single-file .NET exe bundles dozens of framework assemblies,
# each carrying its own version block, so a naive byte scan returns a framework
# version (for example the .NET runtime's) instead of the app's. The only
# reliable value is the RT_VERSION resource reached by walking the PE resource
# directory, which belongs to the outer image alone.
_RT_VERSION = 16

# 0xFEEF04BD little-endian: the dwSignature that opens VS_FIXEDFILEINFO inside
# the RT_VERSION resource data.
_VS_FIXEDFILEINFO_SIGNATURE = b'\xbd\x04\xef\xfe'


def _rva_to_offset(data, rva, sections):
    """Map a resource RVA to a file offset using the section table, or None."""
    for virtual_address, virtual_size, raw_offset, raw_size in sections:
        if virtual_address <= rva < virtual_address + max(virtual_size, raw_size):
            delta = rva - virtual_address
            if delta < raw_size:
                return raw_offset + delta
    return None


def _first_resource_entry_offset(data, table_offset, resource_base, sections,
                                  want_id):
    """Walk one IMAGE_RESOURCE_DIRECTORY level and return a child offset.

    When want_id is not None, follow the entry whose id matches (used to select
    RT_VERSION at the type level); otherwise follow the first entry (used to
    descend name/language levels). Returns the absolute file offset of the child
    directory or data entry, with the high bit stripped, or None.
    """
    # IMAGE_RESOURCE_DIRECTORY: 12 fixed bytes, then NumberOfNamedEntries (2)
    # and NumberOfIdEntries (2), then the entries.
    named = struct.unpack_from('<H', data, table_offset + 12)[0]
    id_count = struct.unpack_from('<H', data, table_offset + 14)[0]
    entries_offset = table_offset + 16
    for index in range(named + id_count):
        entry_offset = entries_offset + index * 8
        name_or_id, offset_to_data = struct.unpack_from('<II', data, entry_offset)
        # Named entries carry the high bit in name_or_id; we only match numeric
        # ids, so skip named entries when selecting by id.
        is_named = bool(name_or_id & 0x80000000)
        if want_id is not None:
            if is_named or name_or_id != want_id:
                continue
        # Directory children have the high bit set on offset_to_data.
        child = (offset_to_data & 0x7FFFFFFF) + resource_base
        return child
    return None


def read_pe_file_version(data):
    """Dotted four-part file version from the PE's RT_VERSION resource, or None.

    Walks the PE resource directory to the RT_VERSION resource of the outer
    executable and decodes its VS_FIXEDFILEINFO file version. Returns None when
    the file has no version resource or the structures are malformed. Never
    raises - callers treat None as "unknown, ask the operator".
    """
    try:
        if len(data) < 0x40 or data[:2] != _PE_MAGIC:
            return None
        pe_offset = struct.unpack_from('<I', data, 0x3C)[0]
        if data[pe_offset:pe_offset + 4] != b'PE\x00\x00':
            return None

        coff = pe_offset + 4
        number_of_sections = struct.unpack_from('<H', data, coff + 2)[0]
        size_of_optional = struct.unpack_from('<H', data, coff + 16)[0]
        optional = coff + 20
        # Magic distinguishes PE32 (0x10b) from PE32+ (0x20b); the data-directory
        # array sits at a different offset in each.
        magic = struct.unpack_from('<H', data, optional)[0]
        if magic == 0x20B:
            data_dir = optional + 112
        else:
            data_dir = optional + 96
        # Resource table is data directory index 2: its RVA then size.
        resource_rva = struct.unpack_from('<I', data, data_dir + 2 * 8)[0]
        if not resource_rva:
            return None

        section_table = optional + size_of_optional
        sections = []
        for index in range(number_of_sections):
            base = section_table + index * 40
            virtual_size = struct.unpack_from('<I', data, base + 8)[0]
            virtual_address = struct.unpack_from('<I', data, base + 12)[0]
            raw_size = struct.unpack_from('<I', data, base + 16)[0]
            raw_offset = struct.unpack_from('<I', data, base + 20)[0]
            sections.append((virtual_address, virtual_size, raw_offset, raw_size))

        resource_base = _rva_to_offset(data, resource_rva, sections)
        if resource_base is None:
            return None

        # Level 1: type -> pick RT_VERSION. Level 2: name/id -> first.
        # Level 3: language -> first. Then a data entry pointing at the resource.
        type_dir = _first_resource_entry_offset(
            data, resource_base, resource_base, sections, _RT_VERSION)
        if type_dir is None:
            return None
        name_dir = _first_resource_entry_offset(
            data, type_dir, resource_base, sections, None)
        if name_dir is None:
            return None
        lang_entry = _first_resource_entry_offset(
            data, name_dir, resource_base, sections, None)
        if lang_entry is None:
            return None

        # lang_entry is an IMAGE_RESOURCE_DATA_ENTRY: RVA (4), size (4), ...
        version_rva = struct.unpack_from('<I', data, lang_entry)[0]
        version_offset = _rva_to_offset(data, version_rva, sections)
        if version_offset is None:
            return None

        # Within the version resource, find VS_FIXEDFILEINFO by its signature.
        signature = data.find(_VS_FIXEDFILEINFO_SIGNATURE, version_offset,
                              version_offset + 512)
        if signature == -1:
            return None
        base = signature + 4 + 4  # skip signature and dwStrucVersion
        version_ms, version_ls = struct.unpack_from('<II', data, base)
        parts = (
            (version_ms >> 16) & 0xFFFF,
            version_ms & 0xFFFF,
            (version_ls >> 16) & 0xFFFF,
            version_ls & 0xFFFF,
        )
        if not any(parts):
            return None
        return '.'.join(str(part) for part in parts)
    except (struct.error, TypeError, IndexError):
        return None


def read_pe_file_version_from_path(path):
    """Read a PE file's version from disk, or None on any failure."""
    try:
        with open(path, 'rb') as handle:
            return read_pe_file_version(handle.read())
    except OSError:
        return None


def _siteagent_dir():
    return os.path.join(current_app.instance_path, 'siteagent')


def _binary_dir(binkey):
    return os.path.join(_siteagent_dir(), binkey)


def _manifest_path(binkey):
    return os.path.join(_binary_dir(binkey), 'manifest.json')


def _legacy_path(binkey):
    return os.path.join(_siteagent_dir(), BINARY_SPECS[binkey]['legacy'])


def _version_exe_path(binkey, version):
    return os.path.join(_binary_dir(binkey), version, BINARY_SPECS[binkey]['exe'])


def _read_manifest(binkey):
    """Return the manifest dict, or an empty skeleton when none is written yet."""
    path = _manifest_path(binkey)
    try:
        with open(path, 'r', encoding='utf-8') as handle:
            data = json.load(handle)
    except (OSError, ValueError):
        return {'current': None, 'versions': []}
    if not isinstance(data, dict):
        return {'current': None, 'versions': []}
    data.setdefault('current', None)
    data.setdefault('versions', [])
    return data


def _write_manifest(binkey, manifest):
    os.makedirs(_binary_dir(binkey), exist_ok=True)
    with open(_manifest_path(binkey), 'w', encoding='utf-8') as handle:
        json.dump(manifest, handle, indent=2)
        handle.write('\n')


def _version_sort_key(version):
    """Parse a dotted-numeric version into a tuple of ints for ordering.

    Pads to four parts so 1.6 and 1.6.0.0 compare equal-ish and shorter
    versions do not sort above longer ones with the same prefix. A part that
    is not numeric (should not happen after validation) sorts as 0.
    """
    parts = []
    for chunk in str(version).split('.'):
        try:
            parts.append(int(chunk))
        except ValueError:
            parts.append(0)
    while len(parts) < 4:
        parts.append(0)
    return tuple(parts[:4])


def list_versions(binkey):
    """Staged versions for a binary, highest version number first.

    Each item is {version, uploadedat, size, iscurrent}. Reads the manifest as
    the source of truth for currency; a version whose exe is missing on disk is
    dropped so the caller never offers a broken push. Ordered by version number
    (descending) so a rollback dropdown reads newest-build-first regardless of
    the order versions were uploaded.
    """
    manifest = _read_manifest(binkey)
    current = manifest.get('current')
    out = []
    for entry in manifest.get('versions', []):
        version = entry.get('version')
        if not version or not os.path.isfile(_version_exe_path(binkey, version)):
            continue
        out.append({
            'version': version,
            'uploadedat': entry.get('uploadedat'),
            'size': entry.get('size'),
            'iscurrent': version == current,
        })
    out.sort(key=lambda item: _version_sort_key(item['version']), reverse=True)
    return out


def current_version(binkey):
    """The current version string, or None when only a legacy flat file exists."""
    manifest = _read_manifest(binkey)
    current = manifest.get('current')
    if current and os.path.isfile(_version_exe_path(binkey, current)):
        return current
    return None


def current_path(binkey):
    """Absolute path to the exe the packager should ship, or None.

    Prefers the current managed version; falls back to the legacy flat file so a
    site that staged a binary before versioning still packages it.
    """
    version = current_version(binkey)
    if version:
        return _version_exe_path(binkey, version)
    legacy = _legacy_path(binkey)
    if os.path.isfile(legacy):
        return legacy
    return None


def legacy_status(binkey):
    """Describe an unmanaged legacy flat file, or None when there is none.

    Returned only when a flat file exists and no managed version is current, so
    the UI can show its detected version and offer one-click adoption without
    the operator re-uploading. {version} is the PE file version if readable,
    else None.
    """
    if current_version(binkey):
        return None
    legacy = _legacy_path(binkey)
    if not os.path.isfile(legacy):
        return None
    try:
        size = os.path.getsize(legacy)
    except OSError:
        size = None
    return {
        'version': read_pe_file_version_from_path(legacy),
        'size': size,
    }


def adopt_legacy(binkey, version):
    """Copy the legacy flat file into the version store as a staged version.

    Reads the flat file, validates it as a PE, and stores it under the given
    version (making it current). Returns the fresh manifest view, or None when
    there is no legacy file or it is not a PE. Caller supplies the version -
    typically the detected PE version, or an operator override when the binary
    carries no version resource.
    """
    legacy = _legacy_path(binkey)
    if not os.path.isfile(legacy):
        return None
    try:
        with open(legacy, 'rb') as handle:
            data = handle.read()
    except OSError:
        return None
    if not looks_like_pe(data):
        return None
    return store_version(binkey, version, data)


def _prune(binkey, manifest):
    """Drop versions past the retention window, never removing the current one."""
    current = manifest.get('current')
    versions = manifest.get('versions', [])
    keep = []
    removed = 0
    for index, entry in enumerate(versions):
        version = entry.get('version')
        # Keep the RETAIN newest (list is newest-first) plus the current always.
        if index < RETAIN or version == current:
            keep.append(entry)
            continue
        # Beyond the window and not current: remove its directory.
        _remove_version_dir(binkey, version)
        removed += 1
    manifest['versions'] = keep
    return removed


def _remove_version_dir(binkey, version):
    directory = os.path.join(_binary_dir(binkey), version)
    try:
        exe = _version_exe_path(binkey, version)
        if os.path.isfile(exe):
            os.remove(exe)
        if os.path.isdir(directory) and not os.listdir(directory):
            os.rmdir(directory)
    except OSError:
        # Best effort: a locked or already-gone file must not fail the upload
        # that triggered pruning. The manifest no longer references it either way.
        pass


def store_version(binkey, version, data, make_current=True):
    """Stage an uploaded binary as a new version.

    Writes data to siteagent/<binkey>/<version>/<exe>, records it newest-first in
    the manifest (de-duping a re-upload of the same version in place), and prunes
    older versions past RETAIN. When make_current is True the new version also
    becomes the current staged build; pass False to stage a rollback/legacy build
    that should appear in the version list without displacing the active current.
    Returns the fresh manifest view {current, versions:[...]}. Caller must have
    validated binkey, version, and that data is a PE binary.
    """
    os.makedirs(os.path.join(_binary_dir(binkey), version), exist_ok=True)
    with open(_version_exe_path(binkey, version), 'wb') as handle:
        handle.write(data)

    manifest = _read_manifest(binkey)
    versions = [entry for entry in manifest.get('versions', [])
                if entry.get('version') != version]
    versions.insert(0, {
        'version': version,
        'uploadedat': datetime.utcnow().isoformat() + 'Z',
        'size': len(data),
        # The identity of the BYTES, not of the label. A version string is
        # chosen by whoever uploads, and re-uploading the same string replaces
        # the file in place - without this, a swap under a trusted label leaves
        # no trace. It is also what an operator compares against the build they
        # meant to ship, and what package time can check before handing an exe
        # to the fleet.
        'sha256': hashlib.sha256(data).hexdigest(),
    })
    manifest['versions'] = versions
    if make_current or not manifest.get('current'):
        manifest['current'] = version
    _prune(binkey, manifest)
    _write_manifest(binkey, manifest)
    return {
        'current': manifest['current'],
        'versions': list_versions(binkey),
    }


def set_current(binkey, version):
    """Point 'current' at an already-staged version. Returns True on success.

    Used to roll back to a retained version. Refuses a version that is not
    staged on disk so current never dangles.
    """
    if not os.path.isfile(_version_exe_path(binkey, version)):
        return False
    manifest = _read_manifest(binkey)
    if not any(entry.get('version') == version
               for entry in manifest.get('versions', [])):
        return False
    manifest['current'] = version
    _write_manifest(binkey, manifest)
    return True


def version_sha256(binkey, version):
    """The recorded digest for a staged version, or None when not recorded.

    Read from the manifest rather than recomputed: store_version already
    hashed the exact bytes it wrote, so re-reading the file to hash it again
    would only introduce a way for the two to disagree.
    """
    for entry in _read_manifest(binkey).get('versions', []):
        if entry.get('version') == version:
            return entry.get('sha256')
    return None


# Name the upgrade payload carries inside the archive. Fixed rather than derived
# from the staged filenames so the install script can address its contents
# without being told what to expect.
PACKAGE_MEMBER_NAMES = {
    'agent': 'shopdbagent.exe',
    'tray': 'shopdb-endpoint-tray.exe',
}

# Fixed member timestamp so packaging the same build twice yields byte-identical
# archives. The value is arbitrary (the DOS epoch zip supports); only its
# constancy matters. Real build provenance lives in the manifest's versions and
# digests, not in a zip entry's mtime.
_PACKAGE_TIMESTAMP = (1980, 1, 1, 0, 0, 0)


def build_current_package():
    """Zip the current agent build (plus the current tray) into one archive.

    Returns (bytes, manifestdict) or (None, None) when no agent build is
    current. The agent is required; the tray is included only when it too has
    a current staged version, because an endpoint can legitimately run the
    service with no tray and half a tray upgrade is worse than none.

    The archive is built IN MEMORY and returned rather than written here: the
    caller decides where it lands, and nothing under the version store gets a
    derived artifact that retention would then have to reason about.

    The returned manifest describes what went in - version and digest per
    member - so the install script can verify it unpacked the build the server
    meant to send, rather than trusting the archive's filename.
    """
    agentversion = current_version('agent')
    if not agentversion:
        return None, None

    members = []
    agentpath = _version_exe_path('agent', agentversion)
    if not os.path.isfile(agentpath):
        return None, None
    members.append(('agent', agentversion, agentpath))

    trayversion = current_version('tray')
    if trayversion:
        traypath = _version_exe_path('tray', trayversion)
        if os.path.isfile(traypath):
            members.append(('tray', trayversion, traypath))

    contents = {}
    buffer = io.BytesIO()
    # ZIP_DEFLATED, not ZIP_STORED: these are ~70 MB self-contained binaries
    # carrying the whole .NET runtime, and the archive crosses the wire to
    # every endpoint being upgraded.
    with zipfile.ZipFile(buffer, 'w', compression=zipfile.ZIP_DEFLATED) as archive:
        for binkey, version, path in members:
            try:
                with open(path, 'rb') as handle:
                    data = handle.read()
            except OSError:
                # A staged file that cannot be read means the package would be
                # silently incomplete. Refuse the whole thing: a partial upgrade
                # payload is how an endpoint ends up with a new service exe and
                # a stale tray.
                return None, None
            name = PACKAGE_MEMBER_NAMES[binkey]
            # Build the entry with a FIXED timestamp. zipfile otherwise stamps
            # each member with the current local time, which makes the archive
            # bytes - and therefore its digest - different on every call for
            # identical input. update_app pins a download to a digest, so a
            # drifting hash would invalidate a package an endpoint is already
            # fetching. Same reason the mode is set explicitly.
            info = zipfile.ZipInfo(name, date_time=_PACKAGE_TIMESTAMP)
            info.compress_type = zipfile.ZIP_DEFLATED
            info.external_attr = 0o644 << 16
            archive.writestr(info, data)
            contents[binkey] = {
                'filename': name,
                'version': version,
                # Prefer the recorded digest so the package agrees with the
                # manifest; fall back to hashing when an older manifest entry
                # predates digest recording.
                'sha256': (version_sha256(binkey, version)
                           or hashlib.sha256(data).hexdigest()),
            }

    payload = buffer.getvalue()
    return payload, {
        'agentversion': agentversion,
        'contents': contents,
        'sha256': hashlib.sha256(payload).hexdigest(),
        'size': len(payload),
    }