#!/usr/bin/env python3
# teledrive_restore.py - Restore TeleDrive files to their original folder structure.
# Copyright (C) 2026 Anders da Silva Rytter Hansen, PC-Rytteren ApS
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, see <https://www.gnu.org/licenses/>.
"""
teledrive_restore.py - Restore TeleDrive files to their original folder structure.

Usage:
    python3 teledrive_restore.py <mv|cp> <files_json> <source_dir> <dest_dir>
    python3 teledrive_restore.py --help
"""

import sys
import json
import os
import shutil

HELP_TEXT = """
teledrive_restore.py - Restore TeleDrive files to their original folder structure

USAGE:
    python3 teledrive_restore.py <operation> <files_json> <source_dir> <dest_dir>

ARGUMENTS:
    operation    Action to perform on each file:
                   mv  - Move files into the restored folder structure
                   cp  - Copy files into the restored folder structure

    files_json   Full path to the TeleDrive export file (typically files.json).
                 The file may have any name as long as it contains valid TeleDrive JSON.

    source_dir   Path to the directory that contains the flat dump of all TeleDrive
                 files. The script will look for each file by its original name inside
                 this directory.

    dest_dir     Path to the destination root directory where the folder tree will be
                 created and files placed. The directory will be created if it does not
                 exist.

BEHAVIOUR:
    - The script reads the TeleDrive metadata and reconstructs the full folder hierarchy
      (including arbitrarily nested sub-folders) under dest_dir.
    - Each file is looked up in source_dir by its original filename (the "name" field in
      the JSON). If the file is not found a WARNING is printed and the script continues.
    - If a destination file already exists it is skipped and a WARNING is printed.
    - Files at the root level (parent_id = null) are placed directly in dest_dir.
    - The script prints a summary of moved/copied, skipped, and missing files when done.

EXAMPLES:
    Copy all files, preserving originals:
        python3 teledrive_restore.py cp /home/user/files.json /mnt/flat_dump /mnt/restored

    Move all files into the restored tree:
        python3 teledrive_restore.py mv /home/user/files.json /mnt/flat_dump /mnt/restored
"""


def build_path_map(entries):
    """Return a dict mapping entry id -> full relative path string."""
    id_map = {e["id"]: e for e in entries}
    path_cache = {}

    def get_path(entry_id):
        if entry_id in path_cache:
            return path_cache[entry_id]
        entry = id_map[entry_id]
        parent_id = entry.get("parent_id")
        if parent_id is None or parent_id not in id_map:
            result = entry["name"]
        else:
            result = os.path.join(get_path(parent_id), entry["name"])
        path_cache[entry_id] = result
        return result

    for e in entries:
        get_path(e["id"])

    return path_cache


def main():
    if len(sys.argv) == 2 and sys.argv[1] in ("--help", "-h"):
        print(HELP_TEXT)
        sys.exit(0)

    if len(sys.argv) != 5:
        print("ERROR: Wrong number of arguments.", file=sys.stderr)
        print("Run with --help for usage information.", file=sys.stderr)
        sys.exit(1)

    operation, json_path, source_dir, dest_dir = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]

    if operation not in ("mv", "cp"):
        print(f"ERROR: First argument must be 'mv' or 'cp', got '{operation}'.", file=sys.stderr)
        sys.exit(1)

    if not os.path.isfile(json_path):
        print(f"ERROR: JSON file not found: {json_path}", file=sys.stderr)
        sys.exit(1)

    if not os.path.isdir(source_dir):
        print(f"ERROR: Source directory not found: {source_dir}", file=sys.stderr)
        sys.exit(1)

    with open(json_path, "r", encoding="utf-8") as f:
        try:
            entries = json.load(f)
        except json.JSONDecodeError as e:
            print(f"ERROR: Failed to parse JSON file: {e}", file=sys.stderr)
            sys.exit(1)

    print(f"Loaded {len(entries)} entries from {json_path}")

    path_map = build_path_map(entries)

    files = [e for e in entries if e.get("type") != "folder"]
    print(f"Files to process: {len(files)}")

    stats = {"done": 0, "missing": 0, "skipped": 0}

    for entry in files:
        filename = entry["name"]
        src = os.path.join(source_dir, filename)

        if not os.path.isfile(src):
            print(f"WARNING: File not found in source, skipping: {filename}")
            stats["missing"] += 1
            continue

        rel_path = path_map[entry["id"]]
        dest_file = os.path.join(dest_dir, rel_path)
        dest_folder = os.path.dirname(dest_file)

        os.makedirs(dest_folder, exist_ok=True)

        if os.path.exists(dest_file):
            print(f"WARNING: Destination already exists, skipping: {dest_file}")
            stats["skipped"] += 1
            continue

        if operation == "cp":
            shutil.copy2(src, dest_file)
        else:
            shutil.move(src, dest_file)

        stats["done"] += 1

    action_word = "Moved" if operation == "mv" else "Copied"
    print()
    print("--- Summary ---")
    print(f"{action_word}: {stats['done']}")
    print(f"Missing (not found in source): {stats['missing']}")
    print(f"Skipped (already at destination): {stats['skipped']}")


if __name__ == "__main__":
    main()