#!/usr/bin/python3

# SPDX-License-Identifier: GPL-2.0-or-later
import grp
import sys
import os
from pathlib import Path
import pwd
import yaml

CONFIG = "/etc/gereal/gereal.conf"
TARGET = "multi-user.target"
DEBUG_MODE = os.getenv("DEBUG") is not None

GENERATOR_DIR = os.environ.get("SYSTEMD_GENERATOR_OUTPUT", "/run/systemd/generator")


def read_config(path):
    """Load and parse configuration file with defaults applied to sessions.

    Reads a YAML configuration file containing 'defaults' and 'sessions' sections.
    Each session in the 'sessions' list is augmented with default values from
    the 'defaults' section for any keys not explicitly defined in the session.
    Sessions are keyed by their 'name' field.

    Args:
        path: Path to the YAML configuration file.

    Returns:
        dict: A dictionary of dictionaries representing sessions, keyed by session
              name, each augmented with default values for missing keys.

    Exits:
        Exits with status 0 if the file is not found or YAML parsing fails.
    """
    if not os.path.isfile(path):
        print(
            f"gereal: config file {path} not found",
            file=sys.stderr,
        )
        sys.exit(0)
    try:
        with open(path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
        defaults = config.get("defaults", {})
        defaults.setdefault("screenrc", "/etc/gereal/screenrc")
        sessions = config.get("sessions", [])
        # Apply defaults to each session and key by name
        augmented_sessions = {}
        for session in sessions:
            # Create a new dict with defaults, then overlay session values
            augmented_session = {**defaults, **session}
            session_name = session.get("name")
            if session_name:
                if "screenrc" not in session:
                    # screenrc is not explicitly set for the session
                    # augmented_session["screenrc" is the default then
                    # if a default session-specific screenrc exists,
                    # overwrite the default
                    per_instance_screenrc = \
                        f"/etc/gereal/screenrc-{session_name}"
                    if os.path.exists(per_instance_screenrc):
                        config["screenrc"] = per_instance_screenrc
                augmented_sessions[session_name] = augmented_session
        return augmented_sessions
    except Exception as e:  # pylint: disable=broad-except
        print(
            f"gereal: failed to parse YAML {path}: {e}",
            file=sys.stderr,
        )
        sys.exit(0)

def user_exists(username):
    try:
        pwd.getpwnam(username)
        return True
    except KeyError:
        return False


def group_exists(groupname):
    try:
        grp.getgrnam(groupname)
        return True
    except KeyError:
        return False


# --- Logging Utility ---
def log(msg):
    tag = "[gereal-generator] "
    if DEBUG_MODE:
        print(tag + msg)


# --- Unit Generation ---
def write_unit(output_dir, target_name, source_path, session):
    name = session.get("name")
    device = session.get("device")
    baudrate = session.get("baudrate")
    user = session.get("user")
    group = session.get("group")

    missing_config = [k for k in ["name", "device", "baudrate", "user", "group"] if not session.get(k)]
    if missing_config:
        print(f"geral-generator: session '{name}' missing fields {', '.join(missing_config)}; skipping", file=sys.stderr)
        return

    missing = []
    if not user_exists(user):
        missing.append(f"user '{user}'")
    if not group_exists(group):
        missing.append(f"group '{group}'")

    if missing:
        # Emit error to stderr; do not generate service
        print(f"gereal-generator: cannot generate unit {name}; missing {', '.join(missing)}", file=sys.stderr)
        sys.exit(0)  # exit 0 is normal for generators, boot continues

    screenrc = "/etc/gereal/screenrc"
    per_instance_screenrc = f"/etc/gereal/screenrc-{name}"
    if os.path.exists(per_instance_screenrc):
        screenrc = per_instance_screenrc
    unit_name = f"gereal-{name}.service"

    unit_content = f"""# generated by gereal.generator
[Unit]
Description=gereal session for {name} on {device}
After={target_name}
SourcePath={source_path}

[Service]
Type=simple
ExecStart=/usr/bin/screen -D -m -L -c "{screenrc}" -s /bin/false -S "gereal:{name}" -t "{name}" "{device}" "{baudrate}"
ExecStartPost=/usr/bin/logger -t gereal-{name} -p daemon.info "attach to session with screen -r gereal:{name} as {user}"
Restart=always
User={user}
Group={group}

[Install]
WantedBy={target_name}
"""

    unit_path = Path(output_dir) / unit_name
    unit_path.write_text(unit_content)
    log(f"Generated unit {unit_name}: {unit_path}")

    return unit_name


def ensure_symlink(target_name, link_dir):
    """
    Create a symlink in `link_dir` pointing to `../target_name`, 
    if it doesn't already exist. Ensures the parent directory exists.
    """
    os.makedirs(link_dir, exist_ok=True)
    log(f"makedirs {link_dir} ok")

    link_path = os.path.join(link_dir, target_name)
    symlink_target = f"../{target_name}"

    if not os.path.exists(link_path):
        log(f"symlink {symlink_target} to {link_path}")
        os.symlink(symlink_target, link_path)
    else:
        log(f"{link_path} already exists")


# --- Entry Point ---
def main():
    if len(sys.argv) != 4:
        print(
            "Usage: generator <unitdir> <runtime-dir> <sysconf-dir>",
            file=sys.stderr
        )
        sys.exit(1)

    output_dir = sys.argv[1]
    runtime_dir = sys.argv[2]
    sysconf_dir = sys.argv[3]
    log(f"Generator started. Output dir: {output_dir}")
    if not Path(CONFIG).exists():
        log(f"Config file not found: {CONFIG}")
        return
    link_target_dir = output_dir + "/" + TARGET + ".wants"

    config = read_config(CONFIG)
    for name,session in config.items():
        log(f"processing {session.get("name", "")}")
        filename = write_unit(output_dir, TARGET, CONFIG, session)
        log(f"generated {filename}")
        if filename and session.get("enable", False):
            ensure_symlink(filename, link_target_dir)
    log("Generator completed.")


if __name__ == "__main__":
    main()
