# import dnf
# import hawkey
import argparse
import json
import tarfile
import os
import shutil
import subprocess
import logging
import shlex
from pathlib import Path
import hashlib
import urllib.request
import pwd

# logging.basicConfig(
#     level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
# )
# logger = logging.getLogger("distrosync")

CMD_RSYNC = "rsync"
CMD_MD5 = "md5"
CMD_APPLY = "apply"
CMD_GENERATE = "generate"
CMD_DIFF = "diff"

REPO_TMP_DIR = "/var/www/repo_tmp"

DIFF_JSON_FILE = "diff.json"
DIFF_TAR_FILE = "diff.tar.gz"

MD5_DEFAULT_NAME = "md5.json"
DOWNLOAD_FOLDER = "download"

REDPESK_USER = "redpesk"
REDPESK_GROUP = "redpesk"

CONTROL_FOLDER_NAME = "_control"

RSYNC_CMD_TEMPLATE = "rsync -avh --prune-empty-dirs\
                                --log-file={log} \
                                --include=imager-os/*** \
                                --include=images/ \
                                --include=images/**/ \
                                --include=images/**/generic/*** \
                                --include=packages/ \
                                --include=packages/**/ \
                                --exclude=*/debug/** \
                                --exclude=*/source/*** \
                                --include=packages/** \
                                --exclude=* \
                                rsync://{url}/redpesk-lts/{distro}-{version}/ \
                                {dest}"


class ColorFormatter(logging.Formatter):
    # ANSI color codes
    blue = "\x1b[34;20m"
    green = "\x1b[32;20m"
    yellow = "\x1b[33;20m"
    red = "\x1b[31;20m"
    bold_red = "\x1b[31;1m"
    reset = "\x1b[0m"

    log_format = "%(asctime)s - %(levelname)s - %(message)s"

    FORMATS = {
        logging.DEBUG: blue + log_format + reset,
        logging.INFO: green + log_format + reset,
        logging.WARNING: yellow + log_format + reset,
        logging.ERROR: red + log_format + reset,
        logging.CRITICAL: bold_red + log_format + reset,
    }

    def format(self, record):
        log_fmt = self.FORMATS.get(record.levelno)
        formatter = logging.Formatter(log_fmt)
        return formatter.format(record)


logger = logging.getLogger("distrosync")
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setFormatter(ColorFormatter())
logger.addHandler(ch)


def initTempFolder(path: str) -> bool:
    """Init a temp folder used before applying a diff.
    This create hard links to current folder.

    Args:
        repo_path (str): repo to use

    Returns:
        bool: True if folder is init, False
    """
    if not os.path.exists(path):
        logger.error(f"Repo {path} not found")
        return False

    logger.info("Creating tmp folder")

    if os.path.exists(REPO_TMP_DIR):
        shutil.rmtree(REPO_TMP_DIR)

    try:
        shutil.copytree(path, REPO_TMP_DIR, copy_function=os.link)

        if os.path.exists(os.path.join(REPO_TMP_DIR, MD5_DEFAULT_NAME)):
            os.remove(os.path.join(REPO_TMP_DIR, MD5_DEFAULT_NAME))
        return True
    except (shutil.Error, Exception) as error:
        logger.error(f"Unable to init tmp folder. Error: {error}")

    return False


def replace_folder(dest: str):
    """Replace current folder by tmp folder after applying diffs.

    Args:
        dest (str): folder to replace
    """
    error = None
    backup = f"{dest.removesuffix('/')}.backup"
    try:
        logger.info(f"Backuping folder {dest} -> {backup}")
        shutil.move(src=dest, dst=backup)

        logger.info(f"Moving {REPO_TMP_DIR} to {dest}")
        shutil.move(src=REPO_TMP_DIR, dst=dest)

        logger.info(f"Removing backup {backup}")
        shutil.rmtree(path=backup)

    except (
        shutil.SameFileError,
        shutil.Error,
        FileNotFoundError,
        PermissionError,
        Exception,
    ) as err:
        logger.error(f"Unable to replace folder. Error: {error}")


def md5sum(filepath):
    md5 = hashlib.md5()

    with open(filepath, "rb") as f:
        while chunk := f.read(1024 * 1024):
            md5.update(chunk)

    return md5.hexdigest()


def download(url: str, output: str) -> bool:
    output = os.path.join(DOWNLOAD_FOLDER, output)
    dir_name = os.path.dirname(output)
    if dir_name and not os.path.exists(dir_name):
        os.makedirs(dir_name, exist_ok=True)

    _download_ok = True
    try:
        logger.info(f"Downloading {url}")
        filename, headers = urllib.request.urlretrieve(url, output)

        if os.path.isfile(filename) and os.path.getsize(filename) < 0:
            logger.error("Download failed: file is empty or missing")
            _download_ok = False

    except urllib.error.HTTPError as e:
        logger.error(f"HTTP error: {e.code} {e.reason}")
        _download_ok = False

    except urllib.error.URLError as e:
        logger.error(f"URL error: {e.reason}")
        _download_ok = False

    except Exception as e:
        logger.error(f"Download failed: {e}")
        _download_ok = False

    return _download_ok


def download_diff_files(url: str, files: list[str], reference_md5: dict) -> bool:
    """Download files from remote. Each downloaded file md5 will be checked.
    If file is already downloaded AND md5 are equals, then skip the download.
    If file is already downloaded BUT md5 are not equals, then warn and download again.

    Args:
        url (str): The base URL to download files from
        files (list[str]): List of file to download, relative to url
        reference_md5 (dict): MD5 file from base url. Used to check downloaded files

    Returns:
        bool: True if OK, False otherwise
    """
    logger.info(f"Downloading {len(files)} files from {url}")
    download_ok = True

    bad_md5 = []

    # create download folder
    if not os.path.exists(DOWNLOAD_FOLDER):
        os.mkdir(DOWNLOAD_FOLDER)

    for f in files:
        _rmd5 = reference_md5.get(f, None)

        if os.path.exists(os.path.join(DOWNLOAD_FOLDER, f)):

            # if already downloaded and MD5 are same, skip. Otherwise, try to download again
            _dmd5 = md5sum(os.path.join(DOWNLOAD_FOLDER, f))
            if _rmd5 == _dmd5:
                logger.debug(f"File {f} already downloaded. Skipping...")
                continue

            # remove previous file if MD5 are not ok
            logger.debug(f"MD5 for File {f} is different [{_dmd5} -> {_rmd5}]. Deleting...")
            os.remove(os.path.join(DOWNLOAD_FOLDER, f))

        download_ok = download(url=f"{url}/{f}", output=f)
        if not download_ok:
            break

        # if downloaded MD5 is not the same as if reference md5 file, warn!!
        _dmd5 = md5sum(os.path.join(DOWNLOAD_FOLDER, f))
        if _rmd5 != _dmd5:
            bad_md5.append({"file": f, "remote_md5": _rmd5, "download_md5": _dmd5})

    # Tell the user to rerun the command
    if bad_md5:
        for bad in bad_md5:
            logger.warning(
                f"Downloaded file {bad.get('file')} has wrong MD5 [{bad.get('remote_md5')} != {bad.get('download_md5')}]"
            )
        logger.warning("=> YOU MUST RERUN THE COMMAND <=")
        return False

    return download_ok


def read_md5_file(path: str) -> dict:
    """Read a MD5 as json file

    Args:
        path (str): Path to the json file containing MD5

    Returns:
        dict: MD5 dict
    """
    if not path or not os.path.exists(path):
        logger.error(f"MD5 file {path} not found")
        return

    logger.info(f"Reading MD5 file {path}")
    md5_data = {}

    try:
        with open(path, "r") as file:
            md5_data = json.loads(file.read())
    except json.decoder.JSONDecodeError as error:
        logger.error(f"Unable to read MD5 file {path}. Error: {error}")
        return

    return md5_data


def create_tar(path: str, files: list[str], reference_path: str, output: str) -> bool:
    """Create a TAR image containing updates.
    The TAR will contains:
    - All files added/downloaded
    - A _control folder used during apply step

    Args:
        path (str): Root folder where files will be found
        files (list[str]): Path to files to add, relative to path
        reference_path (str): Path to the MD5 file used as reference
        output (str): Output path to the generated TAR

    Returns:
        bool: True if OK, False otherwise
    """
    logger.info("Creating TAR")
    diff_path = DIFF_JSON_FILE

    try:
        with tarfile.open(output, "w:gz") as tar:
            # add report
            tar.add(diff_path, arcname=f"{CONTROL_FOLDER_NAME}/{DIFF_JSON_FILE}")

            # add src MD5
            tar.add(reference_path, arcname=f"{CONTROL_FOLDER_NAME}/{MD5_DEFAULT_NAME}")

            for f in files:
                tar.add(os.path.join(path, f), arcname=f)

        logger.info(f"TAR generated: {output}")
        return True

    except FileNotFoundError as error:
        logger.error(error)
        return False


def filter_file(path: str) -> bool:
    """Check if a file must be filtered.

    Args:
        path (str): Path of the file

    Returns:
        bool: False if filtered, True otherwise
    """
    if "debug" in path:
        return False
    if "zephyr" in path:
        return False
    if "source" in path:
        return False
    if "images/smack/minimal/aarch64" in path:
        return False
    if "images/smack/graphic" in path:
        return False

    # following files/folders are use to sync only
    if path.endswith(MD5_DEFAULT_NAME) or "_control" in path:
        return False
    return True


def is_url(path: str) -> bool:
    """Check if a path is a URL. If starts with http:// or https://, this is a URL.

    Args:
        path (str): Path to check

    Returns:
        bool: True if URL, False otherwise
    """
    return path.startswith("http://") or path.startswith("https://")


def check(reference: str) -> bool:
    """Check method used to verify that apply step is OK.

    Args:
        reference (str): The MD5 file used to generate diff

    Returns:
        bool: True if OK, False otherwise
    """
    logger.info("Checking apply")

    reference_md5 = ""
    with open(reference, "r") as file:
        reference_md5 = json.loads(file.read())

    ok = True
    folder = Path(REPO_TMP_DIR)
    for file in folder.rglob("*"):
        if file.is_file():
            if not filter_file(str(file)):
                continue

            _file = str(file).removeprefix(f"{REPO_TMP_DIR}/")

            _md5 = ""
            if _file in reference_md5:
                _md5 = md5sum(file)
                if reference_md5[_file] != _md5:
                    logger.debug(
                        f"File {file} {_md5} must be removed. New MD5 {reference_md5[_file]}"
                    )
                    ok = False
                else:
                    del reference_md5[_file]

            else:
                logger.debug(
                    f"File {_file} {_md5} must be removed. Not found on reference"
                )
                ok = False

    for file in reference_md5:
        logger.debug(f"File {file} must be added")
        ok = False

    if not ok:
        logger.error("Check not ok")
        return False

    logger.info("Check OK")
    return True


def clean_download_folder():
    logger.info(f"Cleaning download folder {DOWNLOAD_FOLDER}")
    if os.path.exists(DOWNLOAD_FOLDER) and os.path.isdir(DOWNLOAD_FOLDER):
        shutil.rmtree(DOWNLOAD_FOLDER)


def change_owner(root_path: str, files: str, user: str | int, group: str | int) -> bool:
    """Change owner for all files/folders added on root_path.

    Args:
        root_path (str): Root path where files belong to
        files (str): Files paths relative to root_path
        user (str | int): Uid
        group (str | int): Gid

    Returns:
        bool: True if OK, False otherwise
    """
    logger.info(f"Changing owner for new files to {user}:{group}")
    try:
        for file in files:
            parent = os.path.join(root_path, file)
            while parent and parent != root_path:
                if os.path.exists(parent):
                    shutil.chown(parent, user, group)
                    logger.debug(f"Owner changed for {parent} -> [{user}:{group}]")
                parent = os.path.dirname(parent)
    except (LookupError, PermissionError, Exception) as error:
        logger.error(f"Unable to change owner for file {file}. Error {error}")
        return False

    return True


def apply(diff_tar: str, dest: str, dry: bool = False):
    """Apply TAR diff based on MD5.

    Args:
        diff_tar (str): tar archive
        dest (str): folder to update
        dry (bool): Dry run. Will not update de dest folder
    """

    logger.info(f"Applying diff from TAR {diff_tar} to folder {dest} [dry: {dry}]")

    if not os.path.exists(diff_tar):
        logger.error(f"TAR file {diff_tar} not found")
        return

    if not os.path.exists(dest):
        logger.error(f"Folder {dest} not found")
        return

    if not initTempFolder(path=dest):
        return

    diff_path = os.path.join(REPO_TMP_DIR, CONTROL_FOLDER_NAME, DIFF_JSON_FILE)

    logger.info(f"Extracting TAR")
    with tarfile.open(diff_tar, "r:gz") as tar:
        tar.extractall(REPO_TMP_DIR, filter="data")

    _diff: dict = None
    with open(diff_path, "r") as f:
        _diff = json.loads(f.read())
    if not _diff:
        logger.error("Unable to read report")
        return

    for removed in _diff.get("removed", []):
        file = os.path.join(REPO_TMP_DIR, str(removed))
        if os.path.isfile(file):
            logger.debug(f"Removing {str(removed)}")
            os.remove(file)

    if not check(
        reference=os.path.join(REPO_TMP_DIR, CONTROL_FOLDER_NAME, MD5_DEFAULT_NAME)
    ):
        return

    if not change_owner(
        root_path=REPO_TMP_DIR,
        files=_diff.get("added", []),
        user=REDPESK_USER,
        group=REDPESK_GROUP,
    ):
        return False

    # check is now ok, so remove _control folder. Not used anymore
    if os.path.exists(os.path.join(REPO_TMP_DIR, CONTROL_FOLDER_NAME)):
        shutil.rmtree(os.path.join(REPO_TMP_DIR, CONTROL_FOLDER_NAME))

    # regenerate MD5 for TMP folder
    if not md5(path=REPO_TMP_DIR, output=os.path.join(REPO_TMP_DIR, MD5_DEFAULT_NAME)):
        logger.error(f"Unable to create new MD5 file for folder {REPO_TMP_DIR}")
        return

    if not dry:
        replace_folder(dest=dest)


def generate(
    path: str,
    output: str = DIFF_TAR_FILE,
    dest_path: str = None,
):
    """Generate TAR file containing updates from a reference folder and an MD5 file.
    if path is an url, then md5 file is not computed but downloaded from this url.
    Otherwise, md5 file is computed based on local folder.

    Args:
        path (str): Path to the folder we want to generate TAR. Could be local or url
        output (str): output tar file.
        dest_path (str): MD5 dest file path. Could be None
    """
    logger.info(
        f"Generating TAR ({output}) diff using folder {path} and {dest_path} MD5"
    )

    reference_path = MD5_DEFAULT_NAME
    if is_url(path):
        if not download(url=f"{path}/{MD5_DEFAULT_NAME}", output=MD5_DEFAULT_NAME):
            return False
        reference_path = os.path.join(DOWNLOAD_FOLDER, MD5_DEFAULT_NAME)
    else:
        # compute MD5 for local folder
        if not md5(path=path, output=MD5_DEFAULT_NAME):
            return False

    diff_dict = diff(reference_path=reference_path, dest_path=dest_path)
    if diff_dict is None:
        return False

    if not diff_dict:
        return True

    if is_url(path):
        # read reference md5 so that we can check downloaded files
        ref_md5 = read_md5_file(reference_path)
        if not ref_md5:
            return False

        if not download_diff_files(
            url=path, files=diff_dict.get("added", []), reference_md5=ref_md5
        ):
            return False
        path = DOWNLOAD_FOLDER

    if not create_tar(
        path=path,
        files=diff_dict.get("added", []),
        reference_path=reference_path,
        output=output,
    ):
        return False

    return True


def md5(path: str, output: str = MD5_DEFAULT_NAME) -> bool:
    """Generate MD5 file for all file in path

    Args:
        path (str): Root path to use
        output (str): output file path
    """
    if not os.path.exists(path) or not os.path.isdir(path):
        logger.error(f"Root path {path} not found")
        return False

    if not path.endswith("/"):
        path = path + "/"

    md5_output = {}

    logger.info(f"Generating MD5 for files in {path}")

    folder = Path(path)
    for file in folder.rglob("*"):
        if file.is_file():
            if filter_file(path=str(file)):
                _f = str(file).removeprefix(path)
                md5_output[_f] = md5sum(file)
                logger.debug(f"MD5: {_f} -> {md5_output[_f]}")

    with open(output, "w") as file:
        file.write(json.dumps(md5_output, indent=4))
    logger.info(f"File generated: {output}")
    return True


def diff(reference_path: str, output_path: str = DIFF_JSON_FILE, dest_path: str = None):
    """Generate diff based on MD5 files. If MD5 files are remote, download them

    Args:
        reference_path (str): MD5 source file path
        output_path (str): Diff output file path
        dest_path (str): MD5 dest file path
    """
    logger.info(
        f"Generating diff between MD5 files reference {reference_path} and {dest_path}"
    )

    # read reference MD5. Download if required
    reference_md5 = None
    reference_path = reference_path
    if is_url(reference_path):
        download(url=reference_path, output=MD5_DEFAULT_NAME)
        reference_path = os.path.join(DOWNLOAD_FOLDER, MD5_DEFAULT_NAME)
    reference_md5 = read_md5_file(path=reference_path)
    if reference_md5 is None:
        return None

    # read dest MD5. Download if required
    dest_md5 = None
    dest_path = dest_path
    if dest_path:
        if is_url(dest_path):
            download(url=dest_path, output="dest.json")
            dest_path = "dest.json"
        dest_md5 = read_md5_file(path=dest_path)
        if dest_md5 is None:
            return None

    # compute diff
    res = {"added": [], "removed": []}
    inter = set(reference_md5.keys()) & set(dest_md5.keys())
    for i in inter:
        if reference_md5[i] != dest_md5[i]:
            res["added"].append(i)
            res["removed"].append(i)
    added = set(reference_md5.keys())
    added.difference_update(inter)
    removed = set(dest_md5.keys())
    removed.difference_update(inter)
    if added:
        res["added"] += list(added)
    if removed:
        res["removed"] += list(removed)

    logger.debug(json.dumps(res, indent=4))

    if not res.get("added") and not res.get("removed"):
        logger.info("No diff detected")
        return {}  # return empty

    # write only if diff detected
    with open(output_path, "w") as file:
        file.write(json.dumps(res, indent=4))
        logger.info(f"Diff generated: {output_path} [{len(added) + len(removed)}]")

    return res


def rsync(url: str, distro: str, version: str):
    """Synchronize a client stack to distro using rsync.
    Network access to distro stack is required.

    Args:
        url (str): Where to download RPM
        distro (str): Name of the distro, i.e batz
        version (str): version of the distro, i.e 2.2-update
    """
    logger.info("RSYNC")
    logger.info(f"Synchronizing {distro}-{version} from {url}")

    _log = f"/distros/logs/{distro}-{version}.log"
    _dest = f"/distros/redpesk-lts/{distro}-{version}"

    if not initTempFolder(_dest):
        return

    cmd = RSYNC_CMD_TEMPLATE.format(
        url=url, distro=distro, version=version, dest=REPO_TMP_DIR, log=_log
    )
    try:
        subprocess.run(shlex.split(cmd), check=True)
        replace_folder(dest_repo=_dest)

    except subprocess.CalledProcessError as _:
        logger.error(f"Error in rsync")
    except KeyboardInterrupt as _:
        logger.error("Rsync stopped")
    finally:
        shutil.rmtree(REPO_TMP_DIR, ignore_errors=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Distro sync for redpesk stacks",
    )
    parser.add_argument(
        "-d", "--debug", help="Enable debug logging", action="store_true", default=False
    )

    subparser = parser.add_subparsers(
        title="Available commands", dest="command", metavar="COMMAND"
    )

    rsync = subparser.add_parser(
        CMD_RSYNC,
        help="Online sync using rsync",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    rsync.add_argument(
        "-u", "--url", help="url to use", default="download.redpesk.bzh/web-share"
    )
    rsync.add_argument(
        "-d",
        "--distro",
        help="distro name to sync",
        required=True,
    )
    rsync.add_argument(
        "-v",
        "--version",
        help="distro version to sync",
        required=True,
    )

    md5_cmd = subparser.add_parser(
        CMD_MD5,
        help="Generate MD5 file for all files in a folder",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    md5_cmd.add_argument(
        "-p",
        "--path",
        help="Path to the folder",
        required=True,
    )
    md5_cmd.add_argument(
        "-o",
        "--output",
        help="Generated MD5 output file",
        required=True,
    )

    generate_cmd = subparser.add_parser(
        CMD_GENERATE,
        help="Generate tar based on MD5 files diff and a reference folder",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    generate_cmd.add_argument(
        "-p",
        "--path",
        help="Path to folder where diff is generated",
        required=True,
    )
    generate_cmd.add_argument(
        "-d",
        "--dest",
        help="Destination MD5 file. MD5 file generated from remote (the folder we want to sync)",
        required=False,
    )
    generate_cmd.add_argument("-o", "--output", help="Output tar path", required=False)
    generate_cmd.add_argument(
        "-k",
        "--keep-download",
        help="Keep download folder. USeful to cache already donwloaded files",
        action="store_true",
        default=False,
        required=False,
    )

    diff_cmd = subparser.add_parser(
        CMD_DIFF,
        help="Generate diff file between 2 MD5 files",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    diff_cmd.add_argument(
        "-s",
        "--src",
        help="MD5 file to use as source",
        required=True,
    )
    diff_cmd.add_argument(
        "-d",
        "--dest",
        help="MD5 file to use as destination",
        required=True,
    )
    diff_cmd.add_argument(
        "-o",
        "--output",
        help=f"Generated diff output file. (Default: {DIFF_JSON_FILE})",
        required=False,
        default=DIFF_JSON_FILE,
    )

    apply_cmd = subparser.add_parser(
        CMD_APPLY,
        help="Apply MD5 diff and copy/remove files",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    apply_cmd.add_argument(
        "-t",
        "--tar",
        help="TAR file to apply",
        required=True,
    )
    apply_cmd.add_argument(
        "-d",
        "--dest",
        help="Destination folder",
        required=True,
    )
    apply_cmd.add_argument(
        "--dry",
        help="Run dry. This will not update the destination folder",
        action="store_true",
        default=False,
        required=False,
    )

    args = parser.parse_args()
    if args.command is None:
        parser.print_help()
        exit(0)

    if args.debug:
        logger.setLevel(logging.DEBUG)

    return args


args = parse_args()

cmd_ret = True
if args.command == CMD_RSYNC:
    rsync(url=args.url, distro=args.distro, version=args.version)
elif args.command == CMD_MD5:
    cmd_ret = md5(path=args.path, output=args.output)
elif args.command == CMD_GENERATE:
    cmd_ret = generate(path=args.path, dest_path=args.dest, output=args.output)

    # clean cache if everything is ok, and TAR has been generated
    if cmd_ret:
        clean_download_folder()

elif args.command == CMD_DIFF:
    cmd_ret = diff(reference=args.src, dest=args.dest, output=args.output)
elif args.command == CMD_APPLY:
    cmd_ret = apply(diff_tar=args.tar, dest=args.dest, dry=args.dry)

if cmd_ret:
    logger.info("Success!")