#!/usr/bin/env python3
"""
ComfyUI WebDAV Output Sync Service

Monitors ComfyUI's output directory and automatically uploads new files
to HiDrive WebDAV storage.

Usage:
    python3 comfyui_webdav_sync.py

Environment Variables:
    WEBDAV_URL: WebDAV server URL (default: https://webdav.hidrive.ionos.com/)
    WEBDAV_USERNAME: WebDAV username
    WEBDAV_PASSWORD: WebDAV password
    WEBDAV_REMOTE_PATH: Remote directory path (default: /users/valknar/Pictures/AI/ComfyUI)
    COMFYUI_OUTPUT_DIR: Local directory to watch (default: /workspace/ComfyUI/output)
    SYNC_DELAY: Seconds to wait after file write before upload (default: 2)
"""

import os
import sys
import time
import logging
from pathlib import Path
from typing import Set
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler, FileCreatedEvent, FileModifiedEvent
from webdav3.client import Client

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger('comfyui-webdav-sync')

# Configuration from environment variables
WEBDAV_URL = os.getenv('WEBDAV_URL')
WEBDAV_USERNAME = os.getenv('WEBDAV_USERNAME')
WEBDAV_PASSWORD = os.getenv('WEBDAV_PASSWORD')
WEBDAV_REMOTE_PATH = os.getenv('WEBDAV_REMOTE_PATH')
COMFYUI_OUTPUT_DIR = os.getenv('COMFYUI_OUTPUT_DIR', '/workspace/ComfyUI/output')
SYNC_DELAY = int(os.getenv('SYNC_DELAY', '2'))  # Wait 2 seconds after file write
RETRY_ATTEMPTS = int(os.getenv('RETRY_ATTEMPTS', '3'))
RETRY_DELAY = int(os.getenv('RETRY_DELAY', '5'))

# Validate required credentials
if not all([WEBDAV_URL, WEBDAV_USERNAME, WEBDAV_PASSWORD, WEBDAV_REMOTE_PATH]):
    logger.error("Missing required WebDAV environment variables!")
    logger.error("Required: WEBDAV_URL, WEBDAV_USERNAME, WEBDAV_PASSWORD, WEBDAV_REMOTE_PATH")
    sys.exit(1)


class WebDAVClient:
    """WebDAV client wrapper with retry logic"""

    def __init__(self):
        options = {
            'webdav_hostname': WEBDAV_URL,
            'webdav_login': WEBDAV_USERNAME,
            'webdav_password': WEBDAV_PASSWORD,
            'webdav_timeout': 300,  # 5 minutes for large files
        }
        self.client = Client(options)
        self._ensure_remote_directory()

    def _ensure_remote_directory(self):
        """Ensure the remote directory exists"""
        try:
            if not self.client.check(WEBDAV_REMOTE_PATH):
                logger.info(f"Creating remote directory: {WEBDAV_REMOTE_PATH}")
                # Create parent directories recursively
                parts = Path(WEBDAV_REMOTE_PATH).parts
                current = ''
                for part in parts:
                    if not part or part == '/':
                        continue
                    current = f"{current}/{part}"
                    if not self.client.check(current):
                        self.client.mkdir(current)
            logger.info(f"✓ Remote directory ready: {WEBDAV_REMOTE_PATH}")
        except Exception as e:
            logger.error(f"Failed to create remote directory: {e}")
            raise

    def upload_file(self, local_path: str, remote_path: str) -> bool:
        """Upload a file with retry logic"""
        for attempt in range(1, RETRY_ATTEMPTS + 1):
            try:
                # Ensure parent directory exists
                remote_dir = str(Path(remote_path).parent)
                if not self.client.check(remote_dir):
                    self.client.mkdir(remote_dir)

                # Upload file
                logger.info(f"[{attempt}/{RETRY_ATTEMPTS}] Uploading {Path(local_path).name} -> {remote_path}")
                self.client.upload_sync(remote_path=remote_path, local_path=local_path)

                # Verify upload
                if self.client.check(remote_path):
                    file_size = os.path.getsize(local_path)
                    logger.info(f"✓ Upload successful: {Path(local_path).name} ({file_size:,} bytes)")
                    return True
                else:
                    logger.warning(f"Upload verification failed for {remote_path}")

            except Exception as e:
                logger.error(f"Upload attempt {attempt} failed: {e}")
                if attempt < RETRY_ATTEMPTS:
                    logger.info(f"Retrying in {RETRY_DELAY} seconds...")
                    time.sleep(RETRY_DELAY)

        logger.error(f"✗ Failed to upload {local_path} after {RETRY_ATTEMPTS} attempts")
        return False


class ComfyUIOutputHandler(FileSystemEventHandler):
    """Handles file system events in ComfyUI output directory"""

    def __init__(self, webdav_client: WebDAVClient):
        self.webdav_client = webdav_client
        self.pending_files: Set[str] = set()  # Files waiting for write completion
        self.uploaded_files: Set[str] = set()  # Track uploaded files
        self.watch_dir = Path(COMFYUI_OUTPUT_DIR)

    def on_created(self, event):
        """Handle file creation events"""
        if event.is_directory:
            return

        file_path = event.src_path

        # Ignore temp files and hidden files
        if self._should_ignore(file_path):
            return

        logger.info(f"New file detected: {Path(file_path).name}")
        self.pending_files.add(file_path)

        # Schedule upload after delay (to ensure file write is complete)
        self._schedule_upload(file_path)

    def on_modified(self, event):
        """Handle file modification events"""
        if event.is_directory:
            return

        file_path = event.src_path

        # Ignore if already uploaded or should be ignored
        if file_path in self.uploaded_files or self._should_ignore(file_path):
            return

        # Add to pending if not already there
        if file_path not in self.pending_files:
            logger.info(f"Modified file detected: {Path(file_path).name}")
            self.pending_files.add(file_path)
            self._schedule_upload(file_path)

    def _should_ignore(self, file_path: str) -> bool:
        """Check if file should be ignored"""
        name = Path(file_path).name

        # Ignore hidden files, temp files, and partial downloads
        if name.startswith('.') or name.endswith('.tmp') or name.endswith('.part'):
            return True

        return False

    def _schedule_upload(self, file_path: str):
        """Schedule file upload after ensuring write is complete"""
        def upload_when_ready():
            time.sleep(SYNC_DELAY)  # Wait for file write to complete

            # Verify file still exists and size is stable
            if not os.path.exists(file_path):
                logger.warning(f"File disappeared: {file_path}")
                self.pending_files.discard(file_path)
                return

            # Check if file size is stable (not being written)
            size1 = os.path.getsize(file_path)
            time.sleep(0.5)
            size2 = os.path.getsize(file_path)

            if size1 != size2:
                logger.info(f"File still being written: {Path(file_path).name}")
                # Reschedule
                self._schedule_upload(file_path)
                return

            # Upload file
            self._upload_file(file_path)

        # Run in background thread
        import threading
        threading.Thread(target=upload_when_ready, daemon=True).start()

    def _upload_file(self, local_path: str):
        """Upload file to WebDAV"""
        try:
            # Calculate relative path from watch directory
            rel_path = Path(local_path).relative_to(self.watch_dir)
            remote_path = f"{WEBDAV_REMOTE_PATH}/{rel_path}".replace('\\', '/')

            # Upload
            success = self.webdav_client.upload_file(local_path, remote_path)

            # Track result
            self.pending_files.discard(local_path)
            if success:
                self.uploaded_files.add(local_path)

        except Exception as e:
            logger.error(f"Failed to upload {local_path}: {e}")
            self.pending_files.discard(local_path)


def main():
    """Main entry point"""
    logger.info("=" * 80)
    logger.info("ComfyUI WebDAV Output Sync Service")
    logger.info("=" * 80)
    logger.info(f"Watch directory: {COMFYUI_OUTPUT_DIR}")
    logger.info(f"WebDAV URL: {WEBDAV_URL}")
    logger.info(f"Remote path: {WEBDAV_REMOTE_PATH}")
    logger.info(f"Sync delay: {SYNC_DELAY}s")
    logger.info("=" * 80)

    # Verify watch directory exists
    if not os.path.exists(COMFYUI_OUTPUT_DIR):
        logger.error(f"Watch directory does not exist: {COMFYUI_OUTPUT_DIR}")
        logger.info(f"Creating watch directory...")
        os.makedirs(COMFYUI_OUTPUT_DIR, exist_ok=True)

    # Initialize WebDAV client
    try:
        webdav_client = WebDAVClient()
    except Exception as e:
        logger.error(f"Failed to initialize WebDAV client: {e}")
        sys.exit(1)

    # Create event handler and observer
    event_handler = ComfyUIOutputHandler(webdav_client)
    observer = Observer()
    observer.schedule(event_handler, COMFYUI_OUTPUT_DIR, recursive=True)

    # Start watching
    observer.start()
    logger.info("✓ Service started - watching for new files...")

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        logger.info("Shutting down...")
        observer.stop()

    observer.join()
    logger.info("Service stopped")


if __name__ == '__main__':
    main()
