#!/usr/bin/env python3
"""
Bidirectional sync between local Obsidian vault and iCloud Drive.

Usage:
    APPLE_EMAIL="user@icloud.com" APPLE_PASSWORD="password" python3 icloud-vault-sync.py

Requires:
    - pip3 install pyicloud
    - Apple ID password (not app-specific) in env vars
    - ~/.pyicloud/ session cache for 2FA skip on repeat runs

Features:
    - Lists iCloud Obsidian vaults
    - Downloads vault contents recursively to local path
    - Uploads local changes back to iCloud
    - Caches created directories to avoid iCloud read-after-write lag
"""

import os
import io
import sys


def load_env_creds():
    """Load Apple credentials from ~/.hermes/.env with *** placeholder handling."""
    env_path = os.path.expanduser("~/.hermes/.env")
    email = os.environ.get("APPLE_EMAIL")
    password = os.environ.get("APPLE_PASSWORD")

    if not email or not password:
        if os.path.exists(env_path):
            with open(env_path, 'r') as f:
                for line in f:
                    line = line.strip()
                    if line.startswith('APPLE_EMAIL=') and not email:
                        email = line.split('=', 1)[1].strip().strip('"').strip("'")
                    elif line.startswith('APPLE_ID_PASSWORD=') and not password:
                        raw = line.split('=', 1)[1].strip().strip('"').strip("'")
                        if raw.startswith('*** '):
                            raw = raw[4:]
                        password = raw
    return email, password


class NamedBytesIO(io.BytesIO):
    def __init__(self, initial_bytes, name):
        super().__init__(initial_bytes)
        self.name = name


def scan_icloud(node, path="", files=None, dirs=None):
    """Recursively scan iCloud Drive node, returning files and dirs maps."""
    if files is None:
        files = {}
    if dirs is None:
        dirs = {}
    children = node.get_children()
    for child in children:
        child_path = f"{path}/{child.name}" if path else child.name
        if child.type == "folder" or child.type == "app_library":
            dirs[child_path] = child
            scan_icloud(child, child_path, files, dirs)
        else:
            files[child_path] = child
    return files, dirs


def get_or_create_dir(parent, dirname, rel_path, cache, known_dirs):
    """Get or create an iCloud directory, caching results."""
    if rel_path in cache:
        return cache[rel_path]
    if rel_path in known_dirs:
        cache[rel_path] = known_dirs[rel_path]
        return known_dirs[rel_path]
    try:
        new_dir = parent.mkdir(dirname)
        cache[rel_path] = new_dir
        return new_dir
    except Exception as e:
        print(f"  [DIR! ] {rel_path} — {e}")
        return None


def sync_down(api, vault_name, local_root):
    """Download an iCloud Obsidian vault to local disk."""
    drive = api.drive
    root = drive.root
    obsidian = root["Obsidian"]
    campaign = obsidian[vault_name]

    icloud_files, icloud_dirs = scan_icloud(campaign)
    print(f"Found {len(icloud_files)} files, {len(icloud_dirs)} dirs on iCloud")

    os.makedirs(local_root, exist_ok=True)

    for rel_path, node in icloud_files.items():
        local_path = os.path.join(local_root, rel_path)
        os.makedirs(os.path.dirname(local_path), exist_ok=True)
        try:
            resp = node.open()
            data = resp.content
            with open(local_path, "wb") as f:
                f.write(data)
            print(f"  [DOWN ] {rel_path} ({len(data)} bytes)")
        except Exception as e:
            print(f"  [ERR  ] {rel_path} — {e}")

    print("Download complete.")


def sync_up(api, vault_name, local_root):
    """Upload local vault changes to iCloud Drive."""
    drive = api.drive
    root = drive.root
    obsidian = root["Obsidian"]
    campaign = obsidian[vault_name]

    icloud_files, icloud_dirs = scan_icloud(campaign)
    print(f"Found {len(icloud_files)} files, {len(icloud_dirs)} dirs on iCloud")

    dir_cache = {}
    uploaded = updated = skipped = errors = 0

    for dirpath, dirnames, filenames in os.walk(local_root):
        rel_dir = os.path.relpath(dirpath, local_root)
        if rel_dir == ".":
            rel_dir = ""

        parent = campaign
        if rel_dir:
            parts = rel_dir.split(os.sep)
            current_path = ""
            for part in parts:
                current_path = f"{current_path}/{part}" if current_path else part
                parent = get_or_create_dir(parent, part, current_path, dir_cache, icloud_dirs)
                if parent is None:
                    break
            if parent is None:
                errors += len(filenames)
                continue

        for filename in filenames:
            local_path = os.path.join(dirpath, filename)
            rel_path = os.path.join(rel_dir, filename) if rel_dir else filename

            with open(local_path, "rb") as f:
                local_data = f.read()

            if rel_path in icloud_files:
                existing = icloud_files[rel_path]
                if existing.size == len(local_data):
                    skipped += 1
                    continue
                else:
                    try:
                        existing.delete()
                    except Exception as e:
                        print(f"  [DEL! ] {rel_path} — {e}")
                        errors += 1
                        continue

            try:
                file_obj = NamedBytesIO(local_data, filename)
                parent.upload(file_obj)
                if rel_path in icloud_files:
                    updated += 1
                    print(f"  [UPD  ] {rel_path} ({len(local_data)} bytes)")
                else:
                    uploaded += 1
                    print(f"  [NEW  ] {rel_path} ({len(local_data)} bytes)")
            except Exception as e:
                print(f"  [ERR  ] {rel_path} — {e}")
                errors += 1

    print(f"\n=== Sync Complete ===")
    print(f"Uploaded: {uploaded}  Updated: {updated}  Skipped: {skipped}  Errors: {errors}")


def main():
    email, password = load_env_creds()
    if not email or not password:
        print("ERROR: Set APPLE_EMAIL and APPLE_PASSWORD env vars or add to ~/.hermes/.env")
        sys.exit(1)

    vault_name = os.environ.get("VAULT_NAME", "Pirate Campain")
    local_root = os.path.expanduser(f"~/obsidian-vault/{vault_name.replace(' ', '-')}")
    mode = os.environ.get("SYNC_MODE", "up")  # 'down' or 'up'

    print(f"Authenticating with Apple ID: {email}")
    from pyicloud import PyiCloudService
    api = PyiCloudService(email, password)

    if api.requires_2fa:
        print("ERROR: 2FA required. Session expired.")
        sys.exit(1)
    print("Login OK")

    if mode == "down":
        sync_down(api, vault_name, local_root)
    else:
        sync_up(api, vault_name, local_root)


if __name__ == "__main__":
    main()