import argparse
import subprocess
import os
import sys
from datetime import datetime

def backup_database(db_url: str, output_dir: str):
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_file = os.path.join(output_dir, f"db_backup_{timestamp}.sql")
    
    print(f"[*] Starting database dump to {backup_file}...")
    cmd = f"pg_dump {db_url} -F p -f {backup_file}"
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    
    if result.returncode == 0:
        print(f"[✓] Backup completed successfully: {backup_file}")
    else:
        print(f"[✗] Backup failed: {result.stderr}")

def restore_database(db_url: str, backup_file: str):
    if not os.path.exists(backup_file):
        print(f"[✗] Backup file not found: {backup_file}")
        return

    print(f"[*] Restoring database from {backup_file}...")
    cmd = f"psql {db_url} -f {backup_file}"
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    
    if result.returncode == 0:
        print("[✓] Restoration completed successfully.")
    else:
        print(f"[✗] Restoration failed: {result.stderr}")

def main():
    parser = argparse.ArgumentParser(description="Database Backup & Restore Utility")
    parser.add_argument("action", choices=["backup", "restore"], help="Action to execute")
    parser.add_argument("--db-url", required=True, help="PostgreSQL connection string")
    parser.add_argument("--dir", default="./backups", help="Backup directory")
    parser.add_argument("--file", help="Backup file path (required for restore)")

    args = parser.parse_args()

    if args.action == "backup":
        backup_database(args.db_url, args.dir)
    elif args.action == "restore":
        if not args.file:
            print("[✗] Error: --file argument is required for restore.")
            sys.exit(1)
        restore_database(args.db_url, args.file)

if __name__ == "__main__":
    main()