from __future__ import absolute_import, print_function

r"""
ComfyUI models sync maintenance script.

Syncs AI models from shared storage to a local drive based on computer name.
- Computers in E_DRIVE_COMPUTERS copy to E:\AI\models
- Computers in D_DRIVE_COMPUTERS copy to D:\AI\models
- All other computers copy to C:\AI\models
Uses the same model list (modellist.txt) for all computers.

Features:
- Computer-specific drive selection
- Delta-aware space checking with 15% safety buffer
- Delta-aware copying (only copies changed/missing files)
- Cleans up files not in the model list
- Uses Robocopy for reliable file operations
- Comprehensive logging for Deadline integration

Designed to be launched on Deadline Workers via the CommandLine or
DeadlineCommand plugin,
so all stdout/stderr (or Deadline logging if available) ends up in job logs.
"""

import os
import ntpath
import re
import shutil
import subprocess
import sys
from datetime import datetime

# Configuration
SOURCE_DIR = os.environ.get("COMFY_MODELS_SOURCE", r"\\YOUR-SERVER\share\AI\models")
MODEL_LIST_PATH = os.environ.get("COMFY_MODEL_LIST", r"\\YOUR-SERVER\share\scripts\modellist.txt")

# Drive configuration by computer
# Add Worker hostnames here for site-specific drive placement.
D_DRIVE_COMPUTERS = set()
E_DRIVE_COMPUTERS = set()
ROBOCOPY_EXECUTABLE = "robocopy"
ROBOCOPY_FLAGS = [
    "/R:10",  # Retry 10 times
    "/W:30",  # Wait 30 seconds between retries
    "/V",     # Verbose output
    "/TS",    # Include source time stamps
    "/FP",    # Include full path names
    "/NP",    # No progress indicator
    "/MT:1"   # Single-threaded for reliability
]

try:
    from Deadline.Scripting import ClientUtils  # type: ignore
except Exception:  # pragma: no cover - Deadline libs unavailable outside Worker
    ClientUtils = None


def log(message):
    """Log to Deadline if available, otherwise stdout."""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = "[ComfyModelsSync] {0} {1}".format(timestamp, message)
    if ClientUtils:
        ClientUtils.LogText(line)
    else:
        print(line)


def get_drive_space_info(dest_dir):
    """Get available space information for destination drive."""
    import shutil

    try:
        # Get drive letter and root path
        drive_letter = dest_dir.split(':')[0].upper()
        drive_root = "{0}:\\".format(drive_letter)

        # Get disk usage stats using cross-platform method
        total_bytes, used_bytes, free_bytes = shutil.disk_usage(drive_root)

        free_gb = free_bytes / (1024**3)
        total_gb = total_bytes / (1024**3)

        return {
            'free_bytes': free_bytes,
            'free_gb': round(free_gb, 2),
            'total_gb': round(total_gb, 2)
        }
    except Exception as e:
        log("WARNING: Could not get drive space info: {0}".format(e))
        return None


def calculate_required_space(valid_model_paths, dest_dir):
    """Calculate space required for files that need copying (delta-aware)."""
    total_delta_bytes = 0
    files_to_copy = []

    log("Analyzing space requirements (delta-aware)...")

    for source_path in valid_model_paths:
        # Calculate destination path
        relative_path = os.path.relpath(checked_model_path(SOURCE_DIR, source_path), SOURCE_DIR)
        dest_path = checked_model_path(dest_dir, os.path.join(dest_dir, relative_path))

        # Check if file needs copying
        needs_copy_flag, _ = needs_copy(source_path, dest_path)

        if needs_copy_flag:
            try:
                file_size = os.path.getsize(source_path)
                total_delta_bytes += file_size
                files_to_copy.append(source_path)
            except OSError:
                log("WARNING: Could not get size for {0}".format(source_path))

    total_delta_gb = round(total_delta_bytes / (1024**3), 2)

    log("Space analysis complete:")
    log("  Files requiring copy: {0}".format(len(files_to_copy)))
    log("  Total space required (delta): {0} GB".format(total_delta_gb))

    return {
        'total_delta_bytes': total_delta_bytes,
        'total_delta_gb': total_delta_gb,
        'files_to_copy': files_to_copy
    }


def get_computer_config():
    """Determine destination drive based on computer name."""
    import socket

    # Get computer name
    computer_name = socket.gethostname().upper()
    log("Running on computer: {0}".format(computer_name))

    # Determine drive
    if computer_name in E_DRIVE_COMPUTERS:
        dest_drive = "E"
    elif computer_name in D_DRIVE_COMPUTERS:
        dest_drive = "D"
    else:
        dest_drive = "C"

    dest_dir = "{0}:\\AI\\models".format(dest_drive)

    log("Using destination drive: {0} (directory: {1})".format(dest_drive, dest_dir))

    return dest_dir


def checked_model_path(root, path):
    if not os.path.isabs(path) or any(c in path for c in '\x00\r\n"'):
        raise ValueError("Model paths must be absolute and contain no control characters or quotes")
    _, tail = ntpath.splitdrive(path)
    if ":" in tail or ".." in tail.replace("\\", "/").split("/") or path.replace("\\", "/").startswith(("//?/", "//./")):
        raise ValueError("Unsafe model path")
    for part in tail.replace("\\", "/").split("/"):
        if part and (part.rstrip(" .") != part or re.fullmatch(r"(?i:CON|PRN|AUX|NUL|COM[1-9]|LPT[1-9])(?:\..*)?", part)):
            raise ValueError("Unsafe model path component")
    root = os.path.normpath(os.path.abspath(root) + os.sep)
    path = os.path.normpath(os.path.abspath(path) + os.sep)
    # Reject foreign shares before resolving links or checking file existence.
    if os.path.commonpath([os.path.normcase(root), os.path.normcase(path)]) != os.path.normcase(root):
        raise ValueError("Model path is outside its configured root")
    real_root = os.path.normcase(os.path.normpath(os.path.realpath(root) + os.sep))
    real_path = os.path.normcase(os.path.normpath(os.path.realpath(path) + os.sep))
    if real_path == real_root or os.path.commonpath([real_root, real_path]) != real_root:
        raise ValueError("Model path escapes its configured root")
    return path


def read_model_list():
    """Read and parse the model list file."""
    if not os.path.exists(MODEL_LIST_PATH):
        raise RuntimeError("Model list file not found: {0}".format(MODEL_LIST_PATH))

    with open(MODEL_LIST_PATH, 'r') as f:
        lines = f.readlines()

    # Parse paths, skip empty lines
    model_paths = []
    for line in lines:
        line = line.strip()
        if line and not line.startswith('#'):
            model_paths.append(checked_model_path(SOURCE_DIR, line))

    log("Loaded {0} model paths from list".format(len(model_paths)))
    return model_paths


def ensure_destination():
    """Ensure destination directory exists."""
    if not os.path.isdir(DEST_DIR):
        log("Destination directory does not exist, creating: {0}".format(DEST_DIR))
        os.makedirs(DEST_DIR)


def get_file_size_mb(file_path):
    """Get file size in MB."""
    try:
        size_bytes = os.path.getsize(file_path)
        return size_bytes / (1024 * 1024)
    except OSError:
        return 0


def analyze_models(model_paths):
    """Analyze model files and return valid/missing lists."""
    valid_files = []
    missing_files = []

    for model_path in model_paths:
        model_path = checked_model_path(SOURCE_DIR, model_path)
        if os.path.exists(model_path):
            valid_files.append(model_path)
        else:
            missing_files.append(model_path)
            log("WARNING: Model file not found: {0}".format(model_path))

    total_size_mb = sum(get_file_size_mb(path) for path in valid_files)

    log("Analysis complete:")
    log("  Total models in list: {0}".format(len(model_paths)))
    log("  Valid source files: {0}".format(len(valid_files)))
    log("  Missing source files: {0}".format(len(missing_files)))
    log("  Total size of valid files: {0:.2f} MB".format(total_size_mb))

    return valid_files, missing_files


def cleanup_extra_files(valid_model_paths):
    """Remove files from destination that are not in the model list."""
    if not os.path.exists(DEST_DIR):
        log("Destination directory does not exist yet - no cleanup needed")
        return 0, 0

    # Build set of expected relative paths
    expected_relative_paths = set()
    for model_path in valid_model_paths:
        relative_path = os.path.relpath(checked_model_path(SOURCE_DIR, model_path), SOURCE_DIR)
        expected_relative_paths.add(relative_path)

    # Find all files in destination
    removed_count = 0
    removed_size_mb = 0

    for root, dirs, files in os.walk(DEST_DIR):
        for file in files:
            full_path = checked_model_path(DEST_DIR, os.path.join(root, file))
            relative_path = os.path.relpath(full_path, DEST_DIR)

            if relative_path not in expected_relative_paths:
                try:
                    size_mb = get_file_size_mb(full_path)
                    os.remove(full_path)
                    log("Removed extra file: {0} ({1:.2f} MB)".format(relative_path, size_mb))
                    removed_count += 1
                    removed_size_mb += size_mb
                except OSError as e:
                    log("WARNING: Failed to remove {0}: {1}".format(relative_path, e))

    # Clean up empty directories
    for root, dirs, files in os.walk(DEST_DIR, topdown=False):
        for dir_name in dirs:
            dir_path = checked_model_path(DEST_DIR, os.path.join(root, dir_name))
            try:
                if not os.listdir(dir_path):
                    os.rmdir(dir_path)
                    log("Removed empty directory: {0}".format(os.path.relpath(dir_path, DEST_DIR)))
            except OSError:
                pass  # Directory not empty or other error

    if removed_count > 0:
        log("Cleanup completed: {0} files removed ({1:.2f} MB freed)".format(removed_count, removed_size_mb))
    else:
        log("Cleanup completed: No extra files found")

    return removed_count, removed_size_mb


def needs_copy(source_path, dest_path):
    """Check if file needs to be copied."""
    if not os.path.exists(dest_path):
        return True, "missing"

    try:
        source_stat = os.stat(source_path)
        dest_stat = os.stat(dest_path)

        if source_stat.st_mtime > dest_stat.st_mtime:
            return True, "newer"
        elif source_stat.st_size != dest_stat.st_size:
            return True, "different size"
        else:
            return False, "up to date"
    except OSError:
        return True, "error checking"


def copy_model_file(source_path, dest_path):
    """Copy a model file using robocopy."""
    source_path = checked_model_path(SOURCE_DIR, source_path)
    dest_path = checked_model_path(DEST_DIR, dest_path)
    source_dir = os.path.dirname(source_path)
    dest_dir = os.path.dirname(dest_path)
    file_name = os.path.basename(source_path)

    # Ensure destination directory exists (including subdirectories)
    os.makedirs(dest_dir, exist_ok=True)

    cmd = [ROBOCOPY_EXECUTABLE, source_dir, dest_dir, file_name] + ROBOCOPY_FLAGS

    log("Running: {0}".format(" ".join(cmd)))
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    stdout, stderr = process.communicate()

    if process.returncode >= 8:
        error_msg = "Robocopy failed with exit code {0}".format(process.returncode)
        if stdout:
            error_msg += "\nSTDOUT: {0}".format(stdout.strip())
        if stderr:
            error_msg += "\nSTDERR: {0}".format(stderr.strip())
        raise RuntimeError(error_msg)

    return process.returncode


def sync_models(valid_model_paths):
    """Sync all valid model files."""
    copied_count = 0
    skipped_count = 0
    error_count = 0
    total_processed = 0

    for source_path in valid_model_paths:
        total_processed += 1

        # Calculate destination path
        relative_path = os.path.relpath(checked_model_path(SOURCE_DIR, source_path), SOURCE_DIR)
        dest_path = checked_model_path(DEST_DIR, os.path.join(DEST_DIR, relative_path))

        # Check if copy is needed
        needs_copy_flag, reason = needs_copy(source_path, dest_path)

        if needs_copy_flag:
            try:
                size_mb = get_file_size_mb(source_path)
                log("Copying ({0}): {1} ({2:.2f} MB)".format(reason, relative_path, size_mb))

                exit_code = copy_model_file(source_path, dest_path)

                if exit_code <= 7:  # Success
                    log("SUCCESS: {0} (exit code: {1})".format(relative_path, exit_code))
                    copied_count += 1
                else:
                    log("WARNING: {0} completed with exit code {1}".format(relative_path, exit_code))
                    copied_count += 1

            except Exception as e:
                log("ERROR: Failed to copy {0}: {1}".format(relative_path, e))
                error_count += 1
        else:
            log("Skipped (up to date): {0}".format(relative_path))
            skipped_count += 1

        # Progress update
        if total_processed % 10 == 0 or total_processed == len(valid_model_paths):
            percent = (total_processed * 100) / len(valid_model_paths)
            log("Progress: {0:.1f}% ({1}/{2}) | Copied: {3} | Skipped: {4} | Errors: {5}".format(
                percent, total_processed, len(valid_model_paths), copied_count, skipped_count, error_count))

    return copied_count, skipped_count, error_count


def main():
    log("ComfyUI models sync starting.")
    log("Source directory: {0}".format(SOURCE_DIR))
    log("Model list: {0}".format(MODEL_LIST_PATH))

    try:
        # Get computer-specific configuration
        dest_dir = get_computer_config()

        # Make dest_dir available globally for other functions
        global DEST_DIR
        DEST_DIR = dest_dir

        # Read model list
        model_paths = read_model_list()

        # Analyze models
        valid_model_paths, missing_files = analyze_models(model_paths)

        # Perform space checking (delta-aware)
        space_info = calculate_required_space(valid_model_paths, dest_dir)
        drive_info = get_drive_space_info(dest_dir)

        if drive_info and space_info:
            # Apply 15% safety buffer like PowerShell script
            buffer_multiplier = 1.15
            required_with_buffer = int(space_info['total_delta_bytes'] * buffer_multiplier)
            required_with_buffer_gb = round(required_with_buffer / (1024**3), 2)

            log("Space check:")
            log("  Required (with 15% buffer): {0} GB".format(required_with_buffer_gb))
            log("  Available: {0} GB".format(drive_info['free_gb']))

            if drive_info['free_bytes'] < required_with_buffer:
                shortfall_gb = round((required_with_buffer - drive_info['free_bytes']) / (1024**3), 2)
                log("ERROR: Insufficient disk space!")
                log("  Shortfall: {0} GB".format(shortfall_gb))
                raise RuntimeError("Insufficient disk space for model sync")
            else:
                log("Space check PASSED")
        else:
            log("WARNING: Could not perform space check - proceeding anyway")

        # Ensure destination exists
        ensure_destination()

        # Cleanup extra files
        cleanup_extra_files(valid_model_paths)

        # Sync models
        copied_count, skipped_count, error_count = sync_models(valid_model_paths)

        # Final summary
        log("Sync completed:")
        log("  Models in list: {0}".format(len(model_paths)))
        log("  Valid source files: {0}".format(len(valid_model_paths)))
        log("  Missing source files: {0}".format(len(missing_files)))
        log("  Files copied: {0}".format(copied_count))
        log("  Files skipped: {0}".format(skipped_count))
        log("  Files with errors: {0}".format(error_count))

        if error_count > 0:
            log("WARNING: Some files failed to copy")
            sys.exit(1)

    except Exception as e:
        log("ERROR: {0}".format(e))
        sys.exit(1)

    log("ComfyUI models sync finished successfully.")


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:  # pragma: no cover - runtime safeguard
        log("ERROR: {0}".format(exc))
        raise


def __main__(*args):  # Deadline's ExecuteScript entry point
    main()