#!/usr/bin/env python3

import argparse
import os
import os.path
import subprocess
import sys


def read_mdt(mdt_file):
    table = []
    try:
        with open(mdt_file, "r") as f:
            keys = None
            for line in f:
                line = line.strip()
                if line.startswith("|-"):
                    continue

                if line.startswith("|"):
                    line = line[1:]

                if line.endswith("|"):
                    line = line[:-1]

                if not keys:
                    keys = [x.strip() for x in line.split("|")]
                    continue

                row = dict()

                for i, cell in enumerate([x.strip() for x in line.split("|")]):
                    if cell != "":
                        row[keys[i]] = cell

                table.append(row)
    except FileNotFoundError:
        pass

    return table


def run_editor(path):
    editor = os.environ.get("EDITOR")
    if not editor:
        editor = os.environ.get("VISUAL", "vim")

    subprocess.check_call([editor, path])


def write_ssh_conf(table):
    output = ""
    for host in sorted(table, key=lambda x: x["Host"]):
        if output != "":
            output += "\n"

        output += f"Host {host['Host']}\n"

        for key in sorted(host.keys()):
            if key in ["Host", "Type"]:
                continue

            output += f"    {key} {host[key]}\n"

    ssh_config_file = os.path.expanduser("~/.ssh/config")

    with open(ssh_config_file, "w") as f:
        f.write(output)

    os.chmod(ssh_config_file, 0o644)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("-n", "--non-interactive", action="store_true", help="")
    group.add_argument("-t", "--edit-types", action="store_true", help="")
    args = parser.parse_args()

    config_folder = os.path.expanduser("~/.config/sshconf/")
    os.makedirs(config_folder, exist_ok=True)

    table_file = f"{config_folder}/sshconf.mdt"
    types_file = f"{config_folder}/types.mdt"

    if not os.path.isfile(types_file):
        with open(types_file, "w") as f:
            f.write("| Type | HostName | Port | User | IdentityFile |\n")
            f.write("|------|----------|------|------|--------------|\n")

    if args.edit_types:
        try:
            run_editor(types_file)
        except subprocess.CalledProcessError:
            sys.exit(1)
        sys.exit(0)

    if not args.non_interactive:
        if not os.path.isfile(table_file):
            with open(table_file, "w") as f:
                f.write("| Host | Type | HostName | Port | User | IdentityFile |\n")
                f.write("|------|------|----------|------|------|--------------|\n")

        try:
            run_editor(table_file)
        except subprocess.CalledProcessError:
            sys.exit(2)

    table = read_mdt(table_file)

    types = dict()
    for line in read_mdt(types_file):
        types[line["Type"]] = line
        del types[line["Type"]]["Type"]

    for host in table:
        host_type = host.get("Type", None)
        if host_type and host_type not in types:
            print(f"Type not found: {host_type}", file=sys.stderr)
            host_type = None

        if host_type:
            for key in types[host_type].keys():
                if key not in host:
                    host[key] = types[host_type][key].replace("<Host>", host["Host"])

        if "HostName" not in host:
            host["HostName"] = host["Host"]

    write_ssh_conf(table)