import time
import logging
import argparse
import sysrsync
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger()

class SyncHandler(FileSystemEventHandler):
    def __init__(self, local_dir, remote_dir, remote_user, remote_host, exclude_dirs, sync_interval, jump_host=None):
        self.local_dir = local_dir
        self.remote_dir = remote_dir
        self.remote_user = remote_user
        self.remote_host = remote_host
        self.exclude_dirs = exclude_dirs
        self.jump_host = jump_host
        self.sync_interval = sync_interval
        self.observer = Observer()

        # Initial sync
        logger.info("Initial sync in both directions")
        self.sync_local_to_remote()
        self.sync_remote_to_local()

        self.observer.schedule(self, path=self.local_dir, recursive=True)
        self.observer.start()

    def on_any_event(self, event):
        if any(excluded in event.src_path for excluded in self.exclude_dirs):
            return
        logger.info(f"Local change detected: {event.src_path} - {event.event_type}")
        self.sync_local_to_remote()

    def sync_local_to_remote(self):
        ssh_command = f"ssh -J {self.jump_host}" if self.jump_host else "ssh"
        try:
            sysrsync.run(
                source=self.local_dir,
                destination=f"{self.remote_user}@{self.remote_host}:{self.remote_dir}",
                options=[
                    "--archive",
                    "--compress",
                    f"-e {ssh_command}"
                ],
                exclusions=self.exclude_dirs
            )
            logger.info("Local => Remote complete")
        except Exception as e:
            logger.error(f"Error during local to remote sync: {e}")

    def sync_remote_to_local(self):
        ssh_command = f"ssh -J {self.jump_host}" if self.jump_host else "ssh"
        try:
            sysrsync.run(
                source=f"{self.remote_user}@{self.remote_host}:{self.remote_dir}",
                destination=self.local_dir,
                options=[
                    "--archive",
                    "--compress",
                    f"-e {ssh_command}"
                ],
                exclusions=self.exclude_dirs
            )
            logger.info("Local <= Remote complete")
        except Exception as e:
            logger.error(f"Error during remote to local sync: {e}")

    def start(self):
        count = 0
        try:
            while True:
                time.sleep(self.sync_interval)
                self.sync_remote_to_local()
                if count % 10 == 0:
                    self.sync_local_to_remote()
                count += 1
        except KeyboardInterrupt:
            self.observer.stop()
        self.observer.join()

def main():
    parser = argparse.ArgumentParser(description="Two-way auto sync script using rsync")
    parser.add_argument("--local-dir", required=True, help="Local directory to sync")
    parser.add_argument("--remote-dir", required=True, help="Remote directory to sync")
    parser.add_argument("--remote-user", default="root", help="Remote user")
    parser.add_argument("--remote-host", required=True, help="Remote host")
    parser.add_argument("--sync-interval", type=int, default=5, help="Sync interval in seconds")
    parser.add_argument("--jump-host", help="Jump host for SSH")
    parser.add_argument("--exclude-dirs", nargs='*', default=[".git", ".venv", "__pycache__", "node_modules", ".DS_Store"], help="Directories to exclude from sync")

    args = parser.parse_args()

    sync_handler = SyncHandler(
        local_dir=args.local_dir,
        remote_dir=args.remote_dir,
        remote_user=args.remote_user,
        remote_host=args.remote_host,
        exclude_dirs=args.exclude_dirs,
        sync_interval=args.sync_interval,
        jump_host=args.jump_host
    )

    # Start the sync handler
    sync_handler.start()

if __name__ == "__main__":
    main()