#!/usr/bin/env python3
"""
iCloud Drive vault lister for Obsidian on Linux.
Loads Apple credentials from ~/.hermes/.env, authenticates with pyicloud,
validates 2FA if needed, and lists iCloud Drive contents.

Usage:
    python3 scripts/icloud-vault-lister.py
    python3 scripts/icloud-vault-lister.py --vault "DND - Joris"
"""
import argparse
import os
from pyicloud import PyiCloudService

ENV_PATH = os.path.expanduser("~/.hermes/.env")


def load_credentials():
    email = None
    password = None
    with open(ENV_PATH, "r") as f:
        for line in f:
            line = line.strip()
            if line.startswith("APPLE_EMAIL="):
                email = line.split("=", 1)[1].strip().strip('"').strip("'")
            elif line.startswith("APPLE_ID_PASSWORD="):
                raw = line.split("=", 1)[1].strip().strip('"').strip("'")
                if raw.startswith("*** "):
                    raw = raw[4:]
                password = raw
    if not email or not password:
        raise RuntimeError("APPLE_EMAIL and APPLE_ID_PASSWORD must be set in ~/.hermes/.env")
    return email, password


def authenticate(email, password):
    api = PyiCloudService(email, password)
    if api.requires_2fa:
        code = input("Apple 2FA code: ")
        if not api.validate_2fa_code(code):
            raise RuntimeError("Invalid 2FA code")
    return api


def list_vaults(api):
    root = api.drive.root
    obsidian = None
    for item in root.get_children():
        if item.type == "app_library" and item.name == "Obsidian":
            obsidian = item
            break
    if not obsidian:
        print("No 'Obsidian' app_library found in iCloud Drive root.")
        return []
    vaults = []
    for vault in obsidian.get_children():
        vaults.append(vault.name)
    return vaults


def list_vault_contents(api, vault_name):
    root = api.drive.root
    obsidian = None
    for item in root.get_children():
        if item.type == "app_library" and item.name == "Obsidian":
            obsidian = item
            break
    if not obsidian:
        raise RuntimeError("Obsidian app_library not found")
    vault = None
    for child in obsidian.get_children():
        if child.name == vault_name:
            vault = child
            break
    if not vault:
        raise RuntimeError(f"Vault '{vault_name}' not found")
    items = []
    for item in vault.get_children():
        items.append((item.type, item.name))
    return items


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="List iCloud Obsidian vaults or a vault's contents.")
    parser.add_argument("--vault", type=str, help="Vault name to list contents of")
    args = parser.parse_args()

    email, password = load_credentials()
    api = authenticate(email, password)

    if args.vault:
        print(f"--- Contents of vault: {args.vault} ---")
        for t, name in list_vault_contents(api, args.vault):
            print(f"  [{t}] {name}")
    else:
        print("--- Available Obsidian vaults ---")
        for v in list_vaults(api):
            print(f"  {v}")