#!/usr/bin/env python3
"""Check or generate AutoMQ Placeholders using a Controller StatefulSet.

Requirements: Python 3.9+, kubectl, PyYAML (python3 -m pip install pyyaml).
Uses the current kubeconfig context unless --context is supplied.
Reads Kubernetes resources only; writes a local manifest only when needed.
Only Placeholders generated by this tool are automatically reused.

Examples:
  python3 check-placeholder.py --instance-id <instance-id>
  kubectl apply -f placeholder.yaml
  python3 check-placeholder.py --instance-id <instance-id>

Exit codes: 0 = all Deployments exist; 1 = missing Deployments; 2 = error.
"""

import argparse
import copy
import hashlib
import json
import re
import shlex
import subprocess
import sys
from decimal import Decimal, InvalidOperation
from pathlib import Path

try:
    import yaml
except ImportError:
    sys.exit('PyYAML is required: python3 -m pip install pyyaml')

MANAGER = 'automq-placeholder-tool'
MANAGED_LABEL = 'app.kubernetes.io/managed-by'
ZONE_KEYS = ('topology.kubernetes.io/zone', 'failure-domain.beta.kubernetes.io/zone')
PRIORITY_CLASS = 'automq-low-priority'
PRIORITY = -1


def digest(value):
    encoded = json.dumps(value, sort_keys=True, separators=(',', ':')).encode()
    return hashlib.sha256(encoded).hexdigest()[:12]


def quantity(value):
    """Normalize Kubernetes CPU/memory quantities for comparisons."""
    match = re.fullmatch(r'([+-]?(?:\d+(?:\.\d*)?|\.\d+))([eE][+-]?\d+|[numkKMGTPE]i?)?', str(value))
    if not match:
        raise ValueError(f'Unsupported resource quantity: {value}')
    number, suffix = match.groups()
    suffix = suffix or ''
    factors = {'': 1, 'n': Decimal('1e-9'), 'u': Decimal('1e-6'), 'm': Decimal('1e-3'),
               'k': 1000, 'K': 1000}
    factors.update({unit: 1000 ** power for power, unit in enumerate('MGTPE', 2)})
    factors.update({unit + 'i': 1024 ** power for power, unit in enumerate('KMGTPE', 1)})
    try:
        exponent = re.fullmatch(r'[eE][+-]?\d+', suffix)
        multiplier = Decimal(10) ** int(suffix[1:]) if exponent else factors[suffix]
        return Decimal(number) * multiplier
    except (InvalidOperation, KeyError, ValueError) as exc:
        raise ValueError(f'Unsupported resource quantity: {value}') from exc


class Kubernetes:
    def __init__(self, context):
        self.command = ['kubectl'] + (['--context', context] if context else [])

    def items(self, resource):
        command = self.command + ['get', resource, '-o', 'json']
        if resource in ('statefulsets', 'deployments'):
            command += ['--all-namespaces']
        if resource == 'statefulsets':
            command += ['-l', 'app.kubernetes.io/component=controller']
        try:
            result = subprocess.run(command, capture_output=True, text=True, timeout=60)
        except (OSError, subprocess.TimeoutExpired) as exc:
            raise ValueError(f'Cannot run kubectl: {exc}') from exc
        if result.returncode:
            raise ValueError(f'{shlex.join(command)} failed:\n{result.stderr.strip()}')
        return json.loads(result.stdout)['items']


def find_controller(statefulsets, instance_id):
    matches = []
    for sts in statefulsets:
        metadata = sts.get('metadata', {})
        labels = {**metadata.get('labels', {}),
                  **sts['spec']['template'].get('metadata', {}).get('labels', {})}
        annotations = metadata.get('annotations', {})
        owned = (labels.get('automq.cloud/instance-id') == instance_id
                 or annotations.get('meta.helm.sh/release-name') == instance_id
                 or metadata.get('name') == f'{instance_id}-automq-enterprise-controller')
        if owned and labels.get('app.kubernetes.io/component') == 'controller':
            matches.append(sts)
    if len(matches) != 1:
        raise ValueError(f'Expected one Controller StatefulSet for {instance_id}; found {len(matches)}. '
                         'Check the Kubernetes context and access to all namespaces.')
    return matches[0]


def controller_config(sts):
    spec = sts['spec']['template']['spec']
    terms = spec.get('affinity', {}).get('nodeAffinity', {}).get(
        'requiredDuringSchedulingIgnoredDuringExecution', {}).get('nodeSelectorTerms', [])
    # CMP emits one AND term. Do not guess how to merge arbitrary OR expressions.
    if len(terms) != 1 or terms[0].get('matchFields'):
        raise ValueError('Expected one nodeAffinity term without matchFields.')
    expressions = copy.deepcopy(terms[0].get('matchExpressions', []))
    for key, value in spec.get('nodeSelector', {}).items():
        expressions.append({'key': key, 'operator': 'In', 'values': [value]})
    values_by_key = {}
    for expression in expressions:
        if expression.get('operator') != 'In' or not expression.get('values'):
            raise ValueError('Node matching conditions must use In with non-empty values.')
        key, values = expression['key'], set(expression['values'])
        values_by_key[key] = values_by_key[key] & values if key in values_by_key else values
        if not values_by_key[key]:
            raise ValueError(f'Node conditions for {key} have no common values.')
    expressions = [{'key': key, 'operator': 'In', 'values': sorted(values)}
                   for key, values in sorted(values_by_key.items())]
    zone_keys = {item['key'] for item in expressions if item['key'] in ZONE_KEYS}
    if len(zone_keys) != 1:
        raise ValueError('Expected one standard Kubernetes zone label in Controller node affinity.')
    zone_key = zone_keys.pop()
    zone_sets = [set(item['values']) for item in expressions if item['key'] == zone_key]
    zones = sorted(set.intersection(*zone_sets))
    if not zones:
        raise ValueError('Controller zone conditions have no common values.')
    containers = [item for item in spec['containers'] if item['name'] == 'automq-controller']
    if len(containers) != 1:
        raise ValueError('Expected an automq-controller container; refusing to infer another container.')
    requests = copy.deepcopy(containers[0].get('resources', {}).get('requests', {}))
    if not {'cpu', 'memory'} <= requests.keys():
        raise ValueError('Controller must specify CPU and memory requests.')
    if any(quantity(value) <= 0 for value in requests.values()):
        raise ValueError('Controller requests must be positive.')
    tolerations_by_value = {}
    for source in spec.get('tolerations', []):
        item = copy.deepcopy(source)
        item.setdefault('operator', 'Equal')
        for key in ('key', 'value', 'effect'):
            if item.get(key) == '':
                item.pop(key)
        if item['operator'] == 'Exists':
            item.pop('value', None)
        tolerations_by_value[json.dumps(item, sort_keys=True)] = item
    tolerations = [tolerations_by_value[key] for key in sorted(tolerations_by_value)]
    config = {'nodeConditions': expressions, 'zoneKey': zone_key,
              'zones': zones, 'tolerations': tolerations,
              'workloadName': sts['spec']['template'].get('metadata', {}).get('labels', {}).get(
                  'app.kubernetes.io/name', 'automq-enterprise'),
              'requests': {key: str(quantity(value).normalize()) for key, value in sorted(requests.items())}}
    return config, requests


def manifests(config, requests, namespace, image, workload_name):
    result = []
    for zone in config['zones']:
        expressions = copy.deepcopy(config['nodeConditions'])
        for expression in expressions:
            if expression['key'] == config['zoneKey']:
                expression['values'] = [zone]
        # Deliberately do not copy the Controller's instance-specific pod affinity.
        anti_affinity = {
            'requiredDuringSchedulingIgnoredDuringExecution': [{
                'namespaceSelector': {},
                'labelSelector': {'matchExpressions': [
                    {'key': 'app.kubernetes.io/name', 'operator': 'In',
                     'values': sorted({'automq-enterprise', workload_name})},
                    {'key': 'app.kubernetes.io/component', 'operator': 'In',
                     'values': ['broker', 'controller']} ]},
                'topologyKey': 'kubernetes.io/hostname'}]}
        # Identity is local to this AZ, independent of the source instance's other AZs.
        fingerprint = {'nodeConditions': expressions, 'tolerations': config['tolerations'],
                       'requests': config['requests'], 'podAntiAffinity': anti_affinity}
        name = f'automq-placeholder-{digest(fingerprint)}'
        selector = {'automq.cloud/placeholder-id': name}
        result.append({
            'apiVersion': 'apps/v1', 'kind': 'Deployment',
            'metadata': {'name': name, 'namespace': namespace,
                         'labels': {MANAGED_LABEL: MANAGER}},
            'spec': {'replicas': 1, 'strategy': {'type': 'Recreate'},
                     'selector': {'matchLabels': selector},
                     'template': {'metadata': {'labels': {
                         **selector, 'app': 'low-priority-placeholder'}}, 'spec': {
                         'automountServiceAccountToken': False,
                         'terminationGracePeriodSeconds': 0,
                         'priorityClassName': PRIORITY_CLASS,
                         'tolerations': config['tolerations'],
                         'affinity': {
                             'nodeAffinity': {'requiredDuringSchedulingIgnoredDuringExecution': {
                                 'nodeSelectorTerms': [{'matchExpressions': expressions}]}},
                             'podAntiAffinity': anti_affinity},
                         'containers': [{'name': 'placeholder', 'image': image,
                                         'resources': {'requests': requests}}]}}}})
    return result


def run(args):
    kube = Kubernetes(args.context)
    sts = find_controller(kube.items('statefulsets'), args.instance_id)
    config, requests = controller_config(sts)
    template = sts['spec']['template']
    classes = {item['metadata']['name']: item for item in kube.items('priorityclasses')}
    if PRIORITY_CLASS in classes and classes[PRIORITY_CLASS]['value'] != PRIORITY:
        raise ValueError(f'{PRIORITY_CLASS} exists with a value other than {PRIORITY}; refusing to change it.')
    workload_name = template.get('metadata', {}).get('labels', {}).get('app.kubernetes.io/name', 'automq-enterprise')
    desired = manifests(config, requests, args.namespace, args.image, workload_name)
    deployments = {(item['metadata']['namespace'], item['metadata']['name']): item
                   for item in kube.items('deployments')}
    missing = []
    print(f"Controller: {sts['metadata']['namespace']}/{sts['metadata']['name']}")
    for zone, deployment in zip(config['zones'], desired):
        metadata = deployment['metadata']
        existing = deployments.get((args.namespace, metadata['name']))
        if existing:
            if existing['metadata'].get('labels', {}).get(MANAGED_LABEL) != MANAGER:
                raise ValueError(f"Name conflict: {args.namespace}/{metadata['name']} is not managed by this tool.")
            print(f'{zone}: EXISTS — {metadata["name"]}')
        else:
            missing.append(deployment)
            print(f'{zone}: MISSING — {metadata["name"]}')
    if missing:
        documents = [{'apiVersion': 'v1', 'kind': 'Namespace', 'metadata': {'name': args.namespace}}]
        if PRIORITY_CLASS not in classes:
            documents.append({'apiVersion': 'scheduling.k8s.io/v1', 'kind': 'PriorityClass',
                              'metadata': {'name': PRIORITY_CLASS}, 'value': PRIORITY,
                              'globalDefault': False, 'description': 'AutoMQ failover Placeholder Pods.'})
        documents.extend(missing)
        output = Path(args.output)
        # Never leave an older, unrelated manifest looking like the current result.
        try:
            stream = output.open('x', encoding='utf-8')
        except FileExistsError as exc:
            raise ValueError(f'Output file already exists: {output}. '
                             'Use --output <new-file.yaml> to choose a new filename.') from exc
        with stream:
            yaml.safe_dump_all(documents, stream, sort_keys=False)
        command = kube.command + ['apply', '-f', str(output)]
        print(f'Generated {output}. Review it, then run:\n  {shlex.join(command)}')
        print('Rerun this script after applying. No Kubernetes resources were changed.')
    else:
        print('All Placeholder Deployments exist. No manifest generated.')
    print('Check Pod status yourself:')
    print('  ' + shlex.join(kube.command + ['get', 'pods', '-n', args.namespace,
                                          '-l', 'app=low-priority-placeholder', '-o', 'wide']))
    return 1 if missing else 0


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--instance-id', required=True)
    parser.add_argument('--context', help='kubectl context (defaults to the current context)')
    parser.add_argument('--namespace', default='automq-placeholders', help='Shared Placeholder namespace')
    parser.add_argument('--output', default='placeholder.yaml', help='New output path; existing files are never overwritten')
    parser.add_argument('--image', default='registry.k8s.io/pause:3.10', help='Placeholder image; override for a private registry')
    args = parser.parse_args()
    if not re.fullmatch(r'[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?', args.namespace):
        parser.error('--namespace must be a valid Kubernetes namespace name')
    try:
        return run(args)
    except (ValueError, OSError, KeyError, TypeError) as exc:
        print(f'ERROR: {exc}', file=sys.stderr)
        return 2


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