#!/usr/bin/python3
"""
Borg List and Restore Helper Script
Lists archives for a backup tag and shows restore examples
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path
from typing import Dict, Any, Optional

try:
    import yaml
except ImportError:
    print("ERROR: PyYAML module is required. Install with: pip3 install PyYAML", file=sys.stderr)
    sys.exit(1)


class BorgListHelper:
    """Helper class to list archives and show restore examples"""

    def __init__(self, config_path: Optional[str] = None):
        self.config = self._load_config(config_path)

    def _load_config(self, config_path: Optional[str] = None) -> Dict[str, Any]:
        """Load configuration from YAML file"""
        config_locations = []

        if config_path:
            config_locations.append(Path(config_path))
        else:
            config_locations.extend([
                Path('/etc/borg/borg.yml'),
                Path.cwd() / 'borg.yml'
            ])

        for config_file in config_locations:
            if config_file.exists():
                try:
                    with open(config_file, 'r') as f:
                        config = yaml.safe_load(f)
                    return config
                except Exception as e:
                    print(f"ERROR: Failed to load config from {config_file}: {e}", file=sys.stderr)
                    sys.exit(1)

        print(f"ERROR: No configuration file found in: {', '.join(str(p) for p in config_locations)}",
              file=sys.stderr)
        sys.exit(1)

    def _get_backup_config(self, tag: str) -> Dict[str, Any]:
        """Get configuration for specific backup tag"""
        backups_config = self.config.get('backups', {})

        if tag not in backups_config:
            print(f"ERROR: Backup tag '{tag}' not found in configuration", file=sys.stderr)
            print(f"Available tags: {', '.join(backups_config.keys())}", file=sys.stderr)
            sys.exit(1)

        return backups_config[tag]

    def list_tags(self):
        """List all available backup tags with their paths"""
        backups_config = self.config.get('backups', {})

        if not backups_config:
            print("No backup tags configured in configuration file.", file=sys.stderr)
            return

        self._print_header("Available Backup Tags")
        print()

        # Calculate max tag length for alignment
        max_tag_len = max(len(tag) for tag in backups_config.keys()) if backups_config else 0

        for tag, config in backups_config.items():
            # Get paths
            paths = config.get('paths', [])
            if isinstance(paths, str):
                paths = [paths]

            # Get repository info
            env = config.get('environment', {})

            # Check if environment is valid
            if not isinstance(env, dict):
                print(f"Tag: {tag:<{max_tag_len}}  [INVALID CONFIG]")
                print(f"  ERROR: 'environment' must be a dictionary in YAML format")
                print(f"  Current format is incorrect. Should be:")
                print(f"    environment:")
                print(f"      BORG_REPO: ssh://user@host:22/./repo")
                print(f"      BORG_PASSPHRASE: your-passphrase")
                print(f"  NOT:")
                print(f"    environment:")
                print(f"      BORG_REPO='...'  # Wrong! Remove '=' and quotes")
                print()
                continue

            repo = env.get('BORG_REPO', 'N/A')

            # Print tag with paths
            print(f"Tag: {tag:<{max_tag_len}}  Paths: {len(paths)}")

            # Print repository
            print(f"  Repository: {repo}")

            # Print paths
            if paths:
                print(f"  Backup paths:")
                for path in paths:
                    print(f"    - {path}")
            else:
                print(f"  Backup paths: (none configured)")

            print()

        print(f"Total tags: {len(backups_config)}")
        print(f"\nUsage: {sys.argv[0]} <tag>")
        print()

    def _get_env_vars(self, backup_config: Dict[str, Any]) -> Dict[str, str]:
        """Prepare environment variables for borg"""
        env = os.environ.copy()

        # Get global borg environment from config (optional)
        global_borg_env = self.config.get('borg', {}).get('environment', {})
        env.update(global_borg_env)

        # Get backup-specific environment variables (these override global settings)
        backup_env = backup_config.get('environment', {})
        env.update(backup_env)

        # Set default BORG_EXIT_CODES if not specified anywhere
        if 'BORG_EXIT_CODES' not in env:
            env['BORG_EXIT_CODES'] = 'modern'

        return env

    def _print_header(self, text: str):
        """Print a formatted header"""
        print("\n" + "=" * 78)
        print(f"  {text}")
        print("=" * 78)

    def _print_export_commands(self, env: Dict[str, str]):
        """Print export commands for manual usage"""
        self._print_header("Required Environment Variables")
        print("\nCopy and paste these commands to use borg manually:\n")

        if 'BORG_REPO' in env:
            print(f"export BORG_REPO='{env['BORG_REPO']}'")
        if 'BORG_PASSPHRASE' in env:
            print(f"export BORG_PASSPHRASE='{env['BORG_PASSPHRASE']}'")
        if 'BORG_EXIT_CODES' in env:
            print(f"export BORG_EXIT_CODES='{env['BORG_EXIT_CODES']}'")

        # Print any additional environment variables
        for key, value in env.items():
            if key.startswith('BORG_') and key not in ['BORG_REPO', 'BORG_PASSPHRASE', 'BORG_EXIT_CODES']:
                print(f"export {key}='{value}'")

        print()

    def _get_latest_archive(self, env: Dict[str, str], last: Optional[int] = None,
                           first: Optional[int] = None, glob_archives: Optional[str] = None) -> Optional[str]:
        """Get the name of the latest archive using the same filters as the main list"""
        try:
            cmd = ['borg', 'list', '--short']

            # Apply the same filters as the main list command
            if last:
                cmd.extend(['--last', str(last)])
            elif first:
                cmd.extend(['--first', str(first)])

            if glob_archives:
                cmd.extend(['--glob-archives', glob_archives])

            result = subprocess.run(
                cmd,
                env=env,
                capture_output=True,
                text=True
            )

            if result.returncode == 0 and result.stdout.strip():
                # Get the last line (equivalent to | tail -n 1)
                archives = result.stdout.strip().split('\n')
                return archives[-1] if archives else None
        except Exception:
            pass
        return None

    def _print_restore_examples(self, archive_template: str, backup_paths: list, env: Dict[str, str],
                               last: Optional[int] = None, first: Optional[int] = None,
                               glob_archives: Optional[str] = None):
        """Print restore examples based on actual backup paths"""
        self._print_header("Restore Examples")

        # Try to get the actual latest archive name using the same filters
        example_archive = self._get_latest_archive(env, last, first, glob_archives)

        # Fallback to generated example if no real archive found
        if not example_archive:
            example_archive = archive_template.replace('{hostname}', os.uname().nodename)
            example_archive = example_archive.replace('{now}', '2025-01-15T02:00:00')
            print("Note: Using example archive name (no actual archives found)\n")

        # Get first backup path for concrete examples
        first_path = backup_paths[0] if backup_paths else '/example/path'
        # Remove trailing slash for cleaner examples
        first_path = first_path.rstrip('/')

        # Extract directory name for examples
        path_name = os.path.basename(first_path) if first_path != '/' else 'root'
        parent_dir = os.path.dirname(first_path) if first_path != '/' else ''

        print("\n=== General Examples ===\n")

        print("1. List all files in a specific archive:")
        print(f"   borg list ::{example_archive}\n")

        print("2. List files matching a pattern in an archive:")
        print(f"   borg list ::{example_archive} --pattern '*.conf'\n")

        print("3. Dry-run extract (show what would be extracted):")
        print(f"   borg extract --dry-run --list ::{example_archive}\n")

        print(f"\n=== Concrete Examples for: {first_path} ===\n")

        print(f"4. List files in the backed up path:")
        if first_path == '/':
            print(f"   borg list ::{example_archive} etc/")
        else:
            print(f"   borg list ::{example_archive} {first_path.lstrip('/')}\n")

        print(f"5. Extract complete path to current directory:")
        print(f"   borg extract ::{example_archive} {first_path.lstrip('/')}\n")
        print(f"   Result: Files will be restored to .{first_path}\n")

        print(f"6. Extract and OVERWRITE existing files (CAUTION!):")
        print(f"   borg extract --force ::{example_archive} {first_path.lstrip('/')}\n")
        print(f"   WARNING: This will overwrite existing files without asking!\n")

        if first_path != '/':
            print(f"7. Restore directly to original location (DANGEROUS - overwrites production!):")
            print(f"   cd /")
            print(f"   borg extract --force ::{example_archive} {first_path.lstrip('/')}\n")
            print(f"   WARNING: This restores directly to {first_path} and overwrites all files!\n")

            print(f"8. Restore to temporary location for inspection:")
            print(f"   mkdir -p /tmp/restore")
            print(f"   cd /tmp/restore")
            print(f"   borg extract ::{example_archive} {first_path.lstrip('/')}\n")
            print(f"   Result: Files will be in /tmp/restore{first_path}\n")

            print(f"9. Restore with stripped path components to custom location:")
            strip_count = first_path.count('/') - 1 if first_path.count('/') > 1 else 1
            print(f"   borg extract --strip-components {strip_count} ::{example_archive} {first_path.lstrip('/')} --target /tmp/restore/{path_name}\n")
            print(f"   Result: Content of {first_path} will be in /tmp/restore/{path_name}\n")

        print(f"10. Extract specific subdirectory:")
        if first_path == '/':
            print(f"   borg extract ::{example_archive} etc/nginx/\n")
        else:
            print(f"   borg extract ::{example_archive} {first_path.lstrip('/')}/subdirectory/\n")

        print(f"11. Extract files matching a pattern from the path:")
        print(f"   borg extract ::{example_archive} --pattern '{first_path.lstrip('/')}/*.conf'\n")

        print("\n=== Safe Restore Workflow (Recommended) ===\n")
        print("1. Create temporary restore directory:")
        print("   mkdir -p /tmp/restore && cd /tmp/restore\n")
        print("2. Extract to temporary location:")
        print(f"   borg extract ::{example_archive} {first_path.lstrip('/')}\n")
        print("3. Verify extracted files:")
        print(f"   ls -la /tmp/restore{first_path}\n")
        print("4. Copy needed files to production (selective restore):")
        print(f"   cp -a /tmp/restore{first_path}/specific-file {first_path}/\n")
        print("5. Or sync entire directory (careful with --delete!):")
        print(f"   rsync -av /tmp/restore{first_path}/ {first_path}/\n")

        print("\nNote: Always test restores in a safe location first!")
        print("      Use 'borg extract --help' for more options.")
        print()

    def list_archives(self, tag: str, show_files: Optional[str] = None,
                      last: Optional[int] = None, first: Optional[int] = None,
                      glob_archives: Optional[str] = None, extra_args: list = None):
        """List archives for a backup tag"""
        backup_config = self._get_backup_config(tag)
        env = self._get_env_vars(backup_config)
        archive_template = backup_config.get('archive_name', '{hostname}-{now}')

        # Get backup paths for concrete examples
        paths = backup_config.get('paths', [])
        if isinstance(paths, str):
            paths = [paths]

        # Print export commands
        self._print_export_commands(env)

        # Build borg list command
        if show_files:
            # List files in a specific archive
            self._print_header(f"Files in Archive: {show_files}")
            cmd = ['borg', 'list', f'::{show_files}']
        else:
            # List all archives
            self._print_header(f"Archives for Tag: {tag}")
            cmd = ['borg', 'list', '--format', '{archive:<50} {time}\n']

            # Add filter options
            if last:
                cmd.extend(['--last', str(last)])
            elif first:
                cmd.extend(['--first', str(first)])

            if glob_archives:
                cmd.extend(['--glob-archives', glob_archives])

        # Add extra arguments if provided
        if extra_args:
            cmd.extend(extra_args)

        # Run borg list
        print(f"\nRunning: {' '.join(cmd)}\n")
        try:
            subprocess.run(cmd, env=env)
        except KeyboardInterrupt:
            print("\n\nInterrupted by user", file=sys.stderr)
            sys.exit(130)
        except Exception as e:
            print(f"\nERROR: Failed to run borg list: {e}", file=sys.stderr)
            sys.exit(1)

        # Print restore examples (using the same filters to get example archive)
        if not show_files:
            self._print_restore_examples(archive_template, paths, env, last, first, glob_archives)


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description='Borg List and Restore Helper - List archives and show restore examples',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # List all available backup tags with their paths
  %(prog)s --list-tags

  # List all archives for backup tag 'minimal'
  %(prog)s minimal

  # List only the last 10 archives
  %(prog)s minimal --last 10

  # List archives from a specific date pattern
  %(prog)s minimal --glob-archives '*2025-01-15*'

  # List only the first 5 archives
  %(prog)s minimal --first 5

  # List files in a specific archive
  %(prog)s minimal --archive server-2025-01-15T02:00:00

  # List archives with custom format
  %(prog)s minimal -- --short

  # Use custom config file
  %(prog)s minimal --config /path/to/borg.yml
        """
    )

    parser.add_argument(
        'tag',
        nargs='?',  # Make tag optional when --list-tags is used
        help='Backup tag to list archives for'
    )

    parser.add_argument(
        '--config',
        '-c',
        help='Path to configuration file (default: /etc/borg/borg.yml or ./borg.yml)',
        default=None
    )

    parser.add_argument(
        '--list-tags',
        action='store_true',
        help='List all available backup tags with their paths and exit',
        default=False
    )

    parser.add_argument(
        '--archive',
        '-a',
        help='Show files in a specific archive (archive name)',
        default=None
    )

    parser.add_argument(
        '--last',
        '-l',
        type=int,
        help='Only show the last N archives',
        default=None
    )

    parser.add_argument(
        '--first',
        '-f',
        type=int,
        help='Only show the first N archives',
        default=None
    )

    parser.add_argument(
        '--glob-archives',
        '-g',
        help='Only show archives matching glob pattern (e.g., "*2025-01-15*")',
        default=None
    )

    parser.add_argument(
        'extra_args',
        nargs='*',
        help='Additional arguments to pass to borg list'
    )

    args = parser.parse_args()

    try:
        helper = BorgListHelper(config_path=args.config)

        # Handle --list-tags option
        if args.list_tags:
            helper.list_tags()
            sys.exit(0)

        # Tag is required when not using --list-tags
        if not args.tag:
            print("ERROR: tag argument is required (or use --list-tags to see available tags)", file=sys.stderr)
            parser.print_help()
            sys.exit(1)

        # Validate mutually exclusive options
        if args.last and args.first:
            print("ERROR: --last and --first are mutually exclusive", file=sys.stderr)
            sys.exit(1)

        helper.list_archives(
            args.tag,
            show_files=args.archive,
            last=args.last,
            first=args.first,
            glob_archives=args.glob_archives,
            extra_args=args.extra_args
        )
    except KeyboardInterrupt:
        print("\nInterrupted by user", file=sys.stderr)
        sys.exit(130)
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()
