#!/opt/imh-python/bin/python3

import argparse
import sys
import subprocess
import os
import pymysql
import re

API_PATH = "/scripts/cwp_api"


def change_primary_domain(username, olddomain, newdomain):
    root_cwp = pymysql.connect(
        read_default_file='/root/.my.cnf', database="root_cwp"
    )
    cur = root_cwp.cursor(pymysql.cursors.DictCursor)

    # Sanity checks

    cur.execute('SELECT * FROM user WHERE username=%s', username)

    result = cur.fetchall()

    if len(result) == 0:
        print("User not found!")
        sys.exit(2)

    if len(result) > 1:
        print("Multipe records for user found!")
        sys.exit(3)

    if result[0]['domain'] != olddomain:
        print(
            f"Account domain {result[0]['domain']} doesn't",
            f"match specified olddomain, {olddomain} ",
        )
        cur.execute('SELECT * FROM user WHERE domain=%s', olddomain)
        result = cur.fetchall()
        if len(result) == 0:
            print(f"Domain {olddomain} not found.")
            sys.exit(4)
        print(f"Domain found on other account {result[0]['username']}")
        sys.exit(5)

    # Update domain in mysql

    cur.execute(
        'UPDATE user SET domain=%s WHERE id=%s', (newdomain, result[0]['id'])
    )
    root_cwp.commit()

    # Rebuild and reload services

    apicmds = ["webservers rebuild_all", "webservers reload"]
    if os.path.isfile(f"/var/named/{olddomain}.db"):
        os.rename(f"/var/named/{olddomain}.db", f"/var/named/{newdomain}.db")
    elif os.path.isfile(f"/var/named/{newdomain}.db"):
        print(
            f"No zone found for {olddomain},",
            f"a zone already exists for {newdomain}",
        )
    else:
        print(
            f"No zone found for {olddomain} or {newdomain},",
            "generating a new zone",
        )
        apicmds.extend([f"account rebuild_var_named {username} {newdomain}"])

    try:
        with open('/etc/named.conf', 'r') as file:
            file_contents = file.read()
    except (FileNotFoundError, PermissionError) as e:
        print(f"Error opening named.conf: {e}")
        sys.exit(1)

    file_contents = re.sub(fr'(^zone.*$)', lambda match: match.group(0).replace(olddomain, newdomain), file_contents, flags=re.MULTILINE)

    try:
        with open('/etc/named.conf', 'w') as file:
            file.write(file_contents)
    except (PermissionError, IOError) as e:
        print(f"Error writing named.conf: {e}")
        sys.exit(1)

    try:
        result = subprocess.run(['named-checkconf', '-z', '/etc/named.conf'], capture_output=True, text=True)
        result.check_returncode()
        print("named.conf check passed.")
    except (subprocess.CalledProcessError, FileNotFoundError) as e:    
        print(f"An error occurred while running 'named-checkconf': {e} \n{e.stdout}")
        sys.exit(1)

    for apicmd in apicmds:
        command = [API_PATH]
        command.extend(apicmd.split(" "))
        print(f"Running {command}")
        subprocess.run(command, check=False)

    subprocess.run(["/opt/imh-cwp-dns/update_zone", newdomain], check=False)

    # Update mail dirs

    if os.path.isdir(f"/var/vmail/{olddomain}"):
        if not os.path.isdir(f"/var/vmail/{newdomain}"):
            print(f"Maildir for {newdomain} found, moving")
            os.rename(f"/var/vmail/{olddomain}", f"/var/vmail/{newdomain}")
        else:
            print(
                f"Target mail dir /var/vmail/{newdomain} already exists!",
                "Please check the mail dirs in /var/vmail",
            )
    else:
        print(f"No maildir for {olddomain} found, not moving")

    sys.exit(0)


parser = argparse.ArgumentParser(description="Change CWP Primary Domain")
parser.add_argument("username", type=str, help="Username to change domain for")
parser.add_argument("olddomain", type=str, help="Old Primary Domain")
parser.add_argument("newdomain", type=str, help="New Primary Domain")


args = parser.parse_args()

change_primary_domain(args.username, args.olddomain, args.newdomain)