#!/usr/bin/python3 -u

# This file originally lived in
# https://github.com/coreos/fedora-coreos-releng-automation. See that repo for
# archeological git research.

'''
    Implements the Fedora CoreOS versioning scheme as per:
        https://github.com/coreos/fedora-coreos-tracker/issues/81
        https://github.com/coreos/fedora-coreos-tracker/issues/211
    And also the RHCOS/SCOS versioning scheme such as:
        9.8.20260125-0
'''

import argparse
import dotenv
import json
import os
import platform
import re
import subprocess
import sys
import time
import yaml

from datetime import datetime

# https://github.com/coreos/fedora-coreos-tracker/issues/211#issuecomment-543547587
FCOS_STREAM_TO_NUM = {
    'next': 1,
    'testing': 2,
    'stable': 3,
    'next-devel': 10,
    'testing-devel': 20,
    'rawhide': 91,
    'branched': 92,
    'bodhi-updates-testing': 93,
    'bodhi-updates': 94,
}

# RHCOS stream to number mapping
RHCOS_STREAM_TO_NUM = {
    'nvidia': '01',
}

def main():
    args = parse_args()
    if args.workdir is not None:
        os.chdir(args.workdir)
    assert os.path.isdir('builds'), 'Missing builds/ dir'

    # Initialize all the components of our versions
    x, y, z, n = (None, None, None, None)

    # Pick up values from our build-args.
    config = dotenv.dotenv_values(args.build_args)

    # Grab the current datetime object representing the timestamp
    # for the timestamp component of our version.
    dt = get_timestamp()

    # The base version in FCOS is a single number (i.e. 43) while
    # in RHCOS/SCOS it's two numbers separated by . (i.e. 9.8 or 10.0)
    x, y = split_base_version(config['VERSION'])

    # For RHCOS with FROM image, extract version from baseimage
    # to preserve the base image's creation so we know from what baseimage
    # the layer was done
    if config['ID'] == 'rhel' and config.get('FROM'):
        base_version = get_version_from_image(config['FROM'])
        if not base_version:
            raise RuntimeError(f"Could not get version from image {config['FROM']}. "
                             "Version label is required for NVIDIA builds.")
        parsed = parse_version(base_version)
        if not parsed:
            raise RuntimeError(f"Could not parse version {base_version} from image {config['FROM']}")
        # Use the base image's version
        x, y, z = parsed[0], parsed[1], parsed[2]
        n = get_next_iteration_nvidia(x, y, z)
    else:
        # The y component in FCOS is the timestamp, while in RHCOS/SCOS
        # it's the z component. We'll convert to a YYYYMMDD formatted string.
        if y is None:
            y = int(dt.strftime('%Y%m%d'))
        else:
            z = int(dt.strftime('%Y%m%d'))

        # At this point if z isn't defined then we're FCOS
        if z is None:
            z = FCOS_STREAM_TO_NUM[config['STREAM']]

        # For !FCOS and in dev mode we'll default to getting the build ID
        # n component by incrementing on top of the last build. For FCOS
        # not in dev mode we'll calculate the n by looking at git history.
        if args.dev or config['ID'] != 'fedora':
            n = get_next_iteration_from_builds(x, y, z)
        else:
            n = get_next_iteration_from_git(str(dt))
    # On !FCOS the delimeter for the `n` component is a -
    n_delimiter = '.' if config['ID'] == 'fedora' else '-'

    # For RHCOS with FROM image
    if config['ID'] == 'rhel' and config.get('FROM'):
        nvidia_prefix = RHCOS_STREAM_TO_NUM['nvidia']
        dev = 'dev' if args.dev else ''
        new_version = f'{x}.{y}.{z}{n_delimiter}{dev}{nvidia_prefix}{n:02d}'
    else:
        # Now we can compute the final version. Note we prepend the
        # `dev` string for the n component if --dev was passed.
        dev = 'dev' if args.dev else ''
        new_version = f'{x}.{y}.{z}{n_delimiter}{dev}{n}'

    eprint(f'VERSIONARY: selected new version for build: {new_version}')

    # sanity check the new version by trying to re-parse it
    parsed = parse_version(new_version)
    eprint(f'VERSIONARY: parse_version returned: {parsed}')
    if parsed is None:
        eprint(f'VERSIONARY ERROR: Failed to parse generated version: {new_version}')
        eprint(f'VERSIONARY ERROR: x={x}, y={y}, z={z}, n={n}, n_delimiter={n_delimiter}')
        sys.exit(1)
    print(new_version)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--build-args', help="path to build-args.conf",
                        default='src/config/build-args.conf')
    parser.add_argument('--workdir', help="path to cosa workdir")
    parser.add_argument(
        "--dev", action="store_true", help="generate a developer version"
    )
    return parser.parse_args()


def get_timestamp():
    """
        Get the timestamp from either the lockfiles, or use the
        current time if no lockfiles exist.
    """

    # XXX: should sanity check that the lockfiles for all the basearches have
    # matching timestamps
    exts = ['json', 'yaml']
    basearch = platform.machine()
    for ext in exts:
        try:
            with open(f"src/config/manifest-lock.{basearch}.{ext}") as f:
                lockfile = yaml.safe_load(f)
                generated = lockfile.get('metadata', {}).get('generated')
                if not generated:
                    raise Exception("Missing 'metadata.generated' key "
                                    f"from {lockfile}")
                dt = datetime.strptime(generated, '%Y-%m-%dT%H:%M:%SZ')
                msg_src = "from lockfile"
                break
        except FileNotFoundError:
            continue
    else:
        msg_src = "from datetime.now()"
        dt = datetime.now()

    eprint(f"timestamp: {dt.strftime('%Y%m%d')} ({msg_src})")
    return dt


def split_base_version(base_version):
    components = base_version.split('.')
    if len(components) == 1:
        return int(components[0]), None
    else:
        return int(components[0]), int(components[1])


def get_version_from_image(from_image):
    """
    Get version from a container image using skopeo inspect.
    Returns the version string or None if not found.
    """
    # Map platform.machine() to container arch names
    machine = platform.machine()
    arch_map = {
        'x86_64': 'amd64',
        'aarch64': 'arm64',
        'arm64': 'arm64',
    }
    arch = arch_map.get(machine, machine)

    try:
        result = subprocess.run(
            [
                'skopeo', 'inspect', '--no-tags',
                '--override-arch', arch,
                '--override-os', 'linux',
                f'docker://{from_image}'
            ],
            capture_output=True,
            text=True,
            check=True
        )
        image_data = json.loads(result.stdout)
        digest = image_data.get('Digest')

        if not digest:
            eprint(f"Warning: Could not find digest for {arch}/linux in manifest list")
            return None

        labels = image_data.get('Labels', {})
        version = labels.get('org.opencontainers.image.version')

        if version:
            # (e.g., "10.2.20260423-0 (Plow)")
            version = version.split()[0]
            return version

    except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError) as e:
        eprint(f"Warning: Could not get version from image: {e}")

    return None


def get_next_iteration_nvidia(x, y, z):
    """
    Get the next iteration number for NVIDIA builds.
    First checks builds/builds.json for existing NVIDIA builds with matching x.y.z.
    If found, increments from the highest build number.
    Otherwise, starts with 01 for the first NVIDIA build of this base version.
    """
    try:
        with open('builds/builds.json') as f:
            builds = json.load(f)
    except FileNotFoundError:
        builds = {'builds': []}

    if len(builds['builds']) == 0:
        return 1

    for build in builds['builds']:
        last_buildid = build['id']
        last_version_tuple = parse_version(last_buildid)
        if not last_version_tuple:
            continue

        # Check if this build matches our x.y.z
        if (x, y, z) == last_version_tuple[:3]:
            # Found a matching build, increment from it
            # Extract just the iteration part (last 2 digits) from nvidia format
            last_n = last_version_tuple[3] % 100  # 101 → 1, 102 → 2, etc.
            n = last_n + 1
            return n

    # No matching builds found, start with 01
    eprint(f"n: 01 (first NVIDIA build for {x}.{y}.{z})")
    return 1


def get_next_iteration_from_builds(x, y, z):
    try:
        with open('builds/builds.json') as f:
            builds = json.load(f)
    except FileNotFoundError:
        builds = {'builds': []}

    if len(builds['builds']) == 0:
        eprint("n: 0 (no previous builds)")
        return 0

    last_buildid = builds['builds'][0]['id']
    last_version_tuple = parse_version(last_buildid)
    if not last_version_tuple:
        eprint(f"n: 0 (previous version {last_buildid} does not match scheme)")
        return 0

    if (x, y, z) != last_version_tuple[:3]:
        eprint(f"n: 0 (previous version {last_buildid} x.y.z does not match)")
        return 0

    n = last_version_tuple[3] + 1
    eprint(f"n: {n} (incremented from previous version {last_buildid})")
    return n


def get_next_iteration_from_git(timestamp):
    """
    Compute the next iteration number based on git commit history.

    Given the Y component of the version (YYYY-MM-DD HH:MM:SS date),
    this counts all commits from the start of that date up to HEAD.
    This guarantees that multiple builds on the same day each receive a
    unique `.n` value, even if several changes occur.

    See: https://github.com/coreos/fedora-coreos-tracker/issues/2015
    """
    try:
        # Count commits after that point
        commit_count_since_change = subprocess.check_output(
            ['git', 'rev-list', '--count', '--after', timestamp, 'HEAD'],
            cwd="src/config", text=True
        ).strip()
        eprint(
            f"n: {commit_count_since_change} "
            "(calculated using git commit history)"
        )
        return int(commit_count_since_change)
    except subprocess.CalledProcessError as err:
        msg = (
            "Git command failed: unable to determine the next "
            f"iteration value ({err})"
        )
        raise RuntimeError(msg) from err


def parse_version(version):
    # Note that (?:pattern) os a non-matching group in python regex so
    # it won't show up in the matched m.groups()
    m = re.match(r'^([0-9]+)\.([0-9]+)\.([0-9]+)(?:\.|-)(?:dev)?([0-9]+)$', version)
    if m is None:
        return None
    # sanity-check date. The time could be in the y component or z component
    # so we have to look for it in either.
    timegroup = 2
    if len(m.group(3)) == 8:
        timegroup = 3
    try:
        time.strptime(m.group(timegroup), '%Y%m%d')
    except ValueError:
        return None

    # For NVIDIA versions, the last group is 4 digits (01xx)
    # Keep the full number (e.g., 0101 → 101, 0102 → 102)
    n = int(m.group(4))

    return (int(m.group(1)), int(m.group(2)), int(m.group(3)), n)


def eprint(*args):
    print(*args, file=sys.stderr)


if __name__ == "__main__":
    sys.exit(main())
