#!/usr/bin/python3
import os
import sys
import subprocess
import re
import time
import ipaddress
import shutil
import tempfile
import xml.etree.ElementTree as ET

# Debian/Ubuntu: /etc/memcached.conf
# Alma:   /etc/sysconfig/memcached (KEY="value" format)
if os.path.exists('/etc/memcached.conf'):
    CONF_PATH = '/etc/memcached.conf'
    CONF_FORMAT = 'flags'   # CLI: -l, -p, -m, -c, -U, -v/-vv/-vvv
elif os.path.exists('/etc/sysconfig/memcached'):
    CONF_PATH = '/etc/sysconfig/memcached'
    CONF_FORMAT = 'sysconfig'  # KEY="value"
else:
    CONF_PATH = '/etc/memcached.conf'
    CONF_FORMAT = 'flags'

SERVICE_NAME = 'memcached'

VERBOSITY_LEVELS = [
    ('0', 'off (no extra output)'),
    ('1', '-v (errors and warnings)'),
    ('2', '-vv (verbose, client commands)'),
    ('3', '-vvv (very verbose, internal state)'),
]

def add_slist(root):
    slist = ET.SubElement(root, 'slist', {'name': 'verbosity'})
    for key, label in VERBOSITY_LEVELS:
        v = ET.SubElement(slist, 'val', {'key': key})
        v.text = label

def parse_config_flags():
    config = {
        'bind': '127.0.0.1',
        'port': '11211',
        'memory': '64',
        'maxconn': '1024',
        'udpport': '0',
        'verbosity': '0',
    }
    if not os.path.exists(CONF_PATH):
        return config

    with open(CONF_PATH, 'r', encoding='utf-8', errors='replace') as f:
        for line in f:
            stripped = line.strip()
            if not stripped or stripped.startswith('#'):
                continue
            
            parts = stripped.split(None, 1)
            flag = parts[0]
            val = parts[1].strip() if len(parts) == 2 else ''

            if flag == '-l':
                config['bind'] = val
            elif flag == '-p':
                config['port'] = val
            elif flag == '-m':
                config['memory'] = val
            elif flag == '-c':
                config['maxconn'] = val
            elif flag == '-U':
                config['udpport'] = val
            elif flag in ('-v', '-vv', '-vvv'):
                config['verbosity'] = str(flag.count('v'))

    return config


def _atomic_write(filepath, lines):
    orig_stat = os.stat(filepath)
    dir_name = os.path.dirname(filepath)
    fd, tmp_path = tempfile.mkstemp(dir=dir_name, suffix='.tmp')
    try:
        os.fchown(fd, orig_stat.st_uid, orig_stat.st_gid)
        os.fchmod(fd, orig_stat.st_mode)
        with os.fdopen(fd, 'w', encoding='utf-8') as f:
            f.writelines(lines)
            f.flush()
            os.fsync(f.fileno())
        os.rename(tmp_path, filepath)
    except Exception:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass
        raise


def backup_config():
    if os.path.exists(CONF_PATH):
        backup = CONF_PATH + '.bak'
        shutil.copy2(CONF_PATH, backup)
        return backup
    return None


def restore_config(backup_path):
    if backup_path and os.path.exists(backup_path):
        shutil.copy2(backup_path, CONF_PATH)
        os.unlink(backup_path)


def write_config_flags(config):
    if not os.path.exists(CONF_PATH):
        return False

    with open(CONF_PATH, 'r', encoding='utf-8', errors='replace') as f:
        lines = f.readlines()

    managed_flags = {'-l', '-p', '-m', '-c', '-U', '-v', '-vv', '-vvv'}
    updated = set()
    new_lines = []

    for line in lines:
        stripped = line.strip()
        if stripped and not stripped.startswith('#'):
            parts = stripped.split(None, 1)
            flag = parts[0]
            if flag in managed_flags:
                if flag in ('-v', '-vv', '-vvv'):
                    if 'verbosity' not in updated:
                        verbosity_line = _verbosity_line(config['verbosity'])
                        if verbosity_line:
                            new_lines.append(verbosity_line + '\n')
                        updated.add('verbosity')
                    continue
                else:
                    key = _flag_to_key(flag)
                    if key and key not in updated:
                        new_lines.append(f"{flag} {config[key]}\n")
                        updated.add(key)
                        continue
                    elif key and key in updated:
                        continue
        new_lines.append(line)

    flag_order = [('-l', 'bind'), ('-p', 'port'), ('-m', 'memory'),
                  ('-c', 'maxconn'), ('-U', 'udpport')]
    for flag, key in flag_order:
        if key not in updated:
            new_lines.append(f"{flag} {config[key]}\n")

    if 'verbosity' not in updated:
        verbosity_line = _verbosity_line(config['verbosity'])
        if verbosity_line:
            new_lines.append(verbosity_line + '\n')

    _atomic_write(CONF_PATH, new_lines)
    return True


def _flag_to_key(flag):
    return {'-l': 'bind', '-p': 'port', '-m': 'memory',
            '-c': 'maxconn', '-U': 'udpport'}.get(flag)


def _verbosity_line(verbosity_val):
    n = int(verbosity_val) if str(verbosity_val).isdigit() else 0
    if n == 0:
        return ''
    return '-' + 'v' * n

def parse_config_sysconfig():
    config = {
        'bind': '127.0.0.1',
        'port': '11211',
        'memory': '64',
        'maxconn': '1024',
        'udpport': '0',
        'verbosity': '0',
    }
    if not os.path.exists(CONF_PATH):
        return config

    raw = {}
    with open(CONF_PATH, 'r', encoding='utf-8', errors='replace') as f:
        for line in f:
            stripped = line.strip()
            if not stripped or stripped.startswith('#'):
                continue
            m = re.match(r'^(\w+)=["\']?([^"\']*)["\']?', stripped)
            if m:
                raw[m.group(1).upper()] = m.group(2).strip()

    if 'PORT' in raw:
        config['port'] = raw['PORT']
    if 'CACHESIZE' in raw:
        config['memory'] = raw['CACHESIZE']
    if 'MAXCONN' in raw:
        config['maxconn'] = raw['MAXCONN']

    options = raw.get('OPTIONS', '')
    m_bind = re.search(r'-l\s+(\S+)', options)
    if m_bind:
        config['bind'] = m_bind.group(1)
    m_udp = re.search(r'-U\s+(\S+)', options)
    if m_udp:
        config['udpport'] = m_udp.group(1)
    v_count = len(re.findall(r'\B-v+\b', options))
    if v_count:
        config['verbosity'] = str(min(v_count, 3))

    return config


def write_config_sysconfig(config):
    if not os.path.exists(CONF_PATH):
        return False

    with open(CONF_PATH, 'r', encoding='utf-8', errors='replace') as f:
        lines = f.readlines()

    options_line_idx = None
    current_options = ''
    for i, line in enumerate(lines):
        m = re.match(r'^OPTIONS=["\']?(.*?)["\']?\s*$', line.strip())
        if m:
            options_line_idx = i
            current_options = m.group(1)

    current_options = re.sub(r'-l\s+\S+', '', current_options)
    current_options = re.sub(r'-U\s+\S+', '', current_options)
    current_options = re.sub(r'\s-v+\b', '', current_options)
    current_options = current_options.strip()

    new_options_parts = [current_options] if current_options else []
    new_options_parts.append(f"-l {config['bind']}")
    new_options_parts.append(f"-U {config['udpport']}")
    verbosity_flag = _verbosity_line(config['verbosity'])
    if verbosity_flag:
        new_options_parts.append(verbosity_flag)
    new_options = ' '.join(new_options_parts).strip()

    updated = set()
    new_lines = []
    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped and not stripped.startswith('#'):
            m = re.match(r'^(\w+)=', stripped)
            if m:
                key = m.group(1).upper()
                if key == 'PORT' and 'PORT' not in updated:
                    new_lines.append(f'PORT="{config["port"]}"\n')
                    updated.add('PORT')
                    continue
                elif key == 'CACHESIZE' and 'CACHESIZE' not in updated:
                    new_lines.append(f'CACHESIZE="{config["memory"]}"\n')
                    updated.add('CACHESIZE')
                    continue
                elif key == 'MAXCONN' and 'MAXCONN' not in updated:
                    new_lines.append(f'MAXCONN="{config["maxconn"]}"\n')
                    updated.add('MAXCONN')
                    continue
                elif key == 'OPTIONS' and 'OPTIONS' not in updated:
                    new_lines.append(f'OPTIONS="{new_options}"\n')
                    updated.add('OPTIONS')
                    continue
                elif key in ('PORT', 'CACHESIZE', 'MAXCONN', 'OPTIONS'):
                    continue  # дубль
        new_lines.append(line)

    if 'PORT' not in updated:
        new_lines.append(f'PORT="{config["port"]}"\n')
    if 'CACHESIZE' not in updated:
        new_lines.append(f'CACHESIZE="{config["memory"]}"\n')
    if 'MAXCONN' not in updated:
        new_lines.append(f'MAXCONN="{config["maxconn"]}"\n')
    if 'OPTIONS' not in updated:
        new_lines.append(f'OPTIONS="{new_options}"\n')

    _atomic_write(CONF_PATH, new_lines)
    return True


def parse_config():
    if CONF_FORMAT == 'sysconfig':
        return parse_config_sysconfig()
    return parse_config_flags()


def write_config(config):
    if CONF_FORMAT == 'sysconfig':
        return write_config_sysconfig(config)
    return write_config_flags(config)

def is_valid_bind(value):
    if not value.strip():
        return False
    for addr in value.split():
        addr = addr.strip()
        if not addr:
            continue
        try:
            ipaddress.ip_address(addr)
        except ValueError:
            return False
    return True


def sanitize(value):
    return value.replace('\n', '').replace('\r', '').replace('\x00', '').strip()


def check_user_level():
    level = int(os.environ.get('AUTH_LEVEL', '0'))
    return level >= 30


root = ET.parse(sys.stdin).getroot()
func = os.environ.get('PARAM_func', '')

if func == 'memcached_custom':
    if not check_user_level():
        err_elem = ET.SubElement(root, 'error')
        ET.SubElement(err_elem, 'msg', {'name': 'body'}).text = 'Access denied'
        ET.dump(root)
        sys.exit(0)

    clicked = os.environ.get('PARAM_clicked_button', '')
    sok = os.environ.get('PARAM_sok', '')

    if clicked == 'save' or sok == 'ok':
        bind = sanitize(os.environ.get('PARAM_listen_addr', '127.0.0.1'))
        port = sanitize(os.environ.get('PARAM_port', '11211'))
        memory = sanitize(os.environ.get('PARAM_memory', '64'))
        maxconn = sanitize(os.environ.get('PARAM_maxconn', '1024'))
        udpport = sanitize(os.environ.get('PARAM_udpport', '0'))
        verbosity = sanitize(os.environ.get('PARAM_verbosity', '0'))

        errors = []
        if not is_valid_bind(bind):
            errors.append('invalid_listen_addr')
        if not re.match(r'^\d+$', port):
            errors.append('invalid_port')
        if not re.match(r'^\d+$', memory):
            errors.append('invalid_memory')
        if not re.match(r'^\d+$', maxconn):
            errors.append('invalid_maxconn')
        if not re.match(r'^\d+$', udpport):
            errors.append('invalid_udpport')
        if verbosity not in ('0', '1', '2', '3'):
            errors.append('invalid_verbosity')

        if errors:
            err_elem = ET.SubElement(root, 'error')
            ET.SubElement(err_elem, 'msg', {'name': 'body'}).text = \
                'Validation failed. Please check numeric fields for invalid characters or letters.'
            for err in errors:
                ET.SubElement(err_elem, 'param', {'name': err.replace('invalid_', '')})
        else:
            config_to_save = {
                'bind': bind,
                'port': port,
                'memory': memory,
                'maxconn': maxconn,
                'udpport': udpport,
                'verbosity': verbosity,
            }

            backup_path = backup_config()
            try:
                if write_config(config_to_save):
                    restart_res = subprocess.run(
                        ['systemctl', 'restart', SERVICE_NAME],
                        capture_output=True, text=True, timeout=15
                    )
                    time.sleep(1)
                    status = subprocess.run(
                        ['systemctl', 'is-active', SERVICE_NAME],
                        capture_output=True, text=True, timeout=5
                    )
                    if status.stdout.strip() != 'active':
                        restore_config(backup_path)
                        subprocess.run(
                            ['systemctl', 'restart', SERVICE_NAME],
                            capture_output=True, text=True, timeout=15
                        )
                        err_elem = ET.SubElement(root, 'error')
                        ET.SubElement(err_elem, 'msg', {'name': 'body'}).text = \
                            'Failed to apply configuration. Changes have been rolled back.'
                        details = (restart_res.stderr or status.stderr or '').strip()
                        if details:
                            ET.SubElement(err_elem, 'msg', {'name': 'hint'}).text = details
                    else:
                        ok_elem = ET.SubElement(root, 'ok', {'msg': 'yes'})
                        ET.SubElement(ok_elem, 'msg', {'name': 'body'}).text = \
                            'Memcached settings saved and restarted.'
                else:
                    err_elem = ET.SubElement(root, 'error')
                    ET.SubElement(err_elem, 'msg', {'name': 'body'}).text = \
                        f'Config file not found: {CONF_PATH}'
            finally:
                if backup_path and os.path.exists(backup_path):
                    os.unlink(backup_path)

    config = parse_config()
    for k, v in config.items():
        xml_key = 'listen_addr' if k == 'bind' else k
        ET.SubElement(root, xml_key).text = v
    add_slist(root)

ET.dump(root)