#!/usr/bin/env python3
"""
IT Pro Partner Portal — Backup & Disaster Recovery

Backs up the portal database + config + uploaded files to S3.
Supports two targets: primary (local region) and cross-region (DR).

Usage:
  python3 backup_portal.py                         # Full backup to configured targets
  python3 backup_portal.py --target primary         # Primary only
  python3 backup_portal.py --target cross_region    # DR only
  python3 backup_portal.py --restore backup_2026-07-03_040001.tar.gz  # Restore from file

Required env vars:
  PORTAL_DB_URL          — sqlite:///path or postgresql://...
  PORTAL_FILES_DIR       — /path/to/uploads
  AWS_ACCESS_KEY_ID      — Wasabi or AWS IAM key
  AWS_SECRET_ACCESS_KEY  — Wasabi or AWS IAM secret

Backup targets (configured in backup_config table or env):
  S3_BUCKET_PRIMARY      — itpropartner-backups-prod
  S3_BUCKET_DR           — itpropartner-backups-dr (cross-region)
  S3_ENDPOINT            — https://s3.us-east-1.wasabisys.com (optional)
"""

import gzip
import hashlib
import io
import json
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

# ─── Configuration ────────────────────────────────────────────────

BACKUP_DIR = Path("/tmp/itpropartner-backups")
RETENTION_DAYS = int(os.environ.get("BACKUP_RETENTION_DAYS", "30"))
TIMESTAMP = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H%M%S")
BACKUP_FILENAME = f"itpropartner_backup_{TIMESTAMP}.tar.gz"
CHECKSUM_FILENAME = f"{BACKUP_FILENAME}.sha256"

# AWS/S3 config
S3_ENDPOINT = os.environ.get("S3_ENDPOINT", "")
S3_REGION = os.environ.get("S3_REGION", "us-east-1")

BACKUP_TARGETS = {}

primary_bucket = os.environ.get("S3_BUCKET_PRIMARY")
if primary_bucket:
    BACKUP_TARGETS["primary"] = {
        "bucket": primary_bucket,
        "region": S3_REGION,
        "endpoint": S3_ENDPOINT,
        "prefix": "portal-backup/",
    }

dr_bucket = os.environ.get("S3_BUCKET_DR")
if dr_bucket:
    BACKUP_TARGETS["cross_region"] = {
        "bucket": dr_bucket,
        "region": os.environ.get("S3_REGION_DR", "us-west-2"),
        "endpoint": os.environ.get("S3_ENDPOINT_DR", ""),
        "prefix": "portal-backup/",
    }


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


# ─── Database ─────────────────────────────────────────────────────

DB_URL = os.environ.get("PORTAL_DB_URL", "")
FILES_DIR = os.environ.get("PORTAL_FILES_DIR", "")


class BackupRunner:
    """Manages the full backup lifecycle for a single target."""

    def __init__(self, target_name: str, target_config: dict):
        self.target_name = target_name
        self.bucket = target_config["bucket"]
        self.region = target_config["region"]
        self.endpoint = target_config["endpoint"]
        self.prefix = target_config["prefix"]
        self.run_id = str(uuid.uuid4())
        self.start_time = time.time()
        self.status = "running"
        self.error = None
        self.db_size = 0
        self.files_size = 0
        self.archive_size = 0
        self.checksum = ""

    def run(self) -> dict:
        """Execute backup: dump DB, collect files, archive, upload."""
        try:
            BACKUP_DIR.mkdir(parents=True, exist_ok=True)
            work_dir = BACKUP_DIR / self.target_name
            if work_dir.exists():
                shutil.rmtree(work_dir)
            work_dir.mkdir(parents=True)

            # 1. Dump database
            db_path = self._dump_database(work_dir)
            if db_path:
                self.db_size = os.path.getsize(db_path)

            # 2. Collect files (config, uploads, scripts)
            files_path = self._collect_files(work_dir)

            # 3. Write metadata
            self._write_metadata(work_dir)

            # 4. Create compressed archive
            archive_path, self.archive_size, self.checksum = self._archive(work_dir)

            # 5. Upload to S3
            self._upload(archive_path)

            # 6. Cleanup old backups
            self._cleanup_old()

            # 7. Cleanup temp
            shutil.rmtree(work_dir)

            self.status = "success"
        except Exception as e:
            self.status = "failed"
            self.error = str(e)
            raise

        finally:
            self.duration = time.time() - self.start_time

        return {
            "target": self.target_name,
            "bucket": self.bucket,
            "status": self.status,
            "archive_key": f"{self.prefix}{BACKUP_FILENAME}",
            "checksum": self.checksum,
            "size_bytes": self.archive_size,
            "duration_seconds": round(self.duration, 2),
            "db_backup_size": self.db_size,
            "files_backup_size": self.files_size,
        }

    def _dump_database(self, work_dir: Path) -> Path | None:
        """Dump the database to a compressed SQL file."""
        if not DB_URL:
            return None

        db_file = work_dir / "portal-database.sql"
        db_gz_file = work_dir / "portal-database.sql.gz"

        print(f"  📦 Dumping database: {DB_URL[:60]}...")

        if DB_URL.startswith("sqlite://"):
            # SQLite — just copy the file
            sqlite_path = DB_URL.replace("sqlite:///", "")
            if not sqlite_path.startswith("/"):
                sqlite_path = os.path.expanduser(f"~/{sqlite_path}")
            if not os.path.exists(sqlite_path):
                print(f"  ⚠️  SQLite file not found: {sqlite_path}")
                return None

            # Dump to SQL
            with open(sqlite_path) as src:
                data = src.read()
            with gzip.open(db_gz_file, "wt") as dst:
                dst.write(f"-- SQLite dump from {sqlite_path}\n-- Generated: {now_iso()}\n\n")
                dst.write(data)
            return db_gz_file

        elif DB_URL.startswith("postgresql"):
            # PostgreSQL — use pg_dump
            try:
                result = subprocess.run(
                    ["pg_dump", "--no-owner", "--no-acl", "--compress=9", "--file", str(db_gz_file), DB_URL],
                    capture_output=True, text=True, timeout=120,
                )
                if result.returncode != 0:
                    print(f"  ⚠️  pg_dump failed: {result.stderr[:200]}")
                    return None
                return db_gz_file
            except FileNotFoundError:
                print("  ⚠️  pg_dump not installed")
                return None
            except subprocess.TimeoutExpired:
                print("  ⚠️  pg_dump timed out")
                return None

        print(f"  ⚠️  Unsupported DB URL: {DB_URL.split('://')[0]}://...")
        return None

    def _collect_files(self, work_dir: Path) -> Path:
        """Collect config files, uploads, and scripts."""
        files_dir = work_dir / "files"
        files_dir.mkdir()

        # Portal config directory
        config_dirs = [
            FILES_DIR if FILES_DIR and os.path.exists(FILES_DIR) else None,
            Path(os.path.expanduser("~/.hermes")),
        ]

        for src_dir in config_dirs:
            if not src_dir or not Path(src_dir).exists():
                continue
            src_path = Path(src_dir)
            dst_path = files_dir / src_path.name
            try:
                shutil.copytree(src_path, dst_path, ignore=shutil.ignore_patterns("*.pyc", "__pycache__", ".git"))
                size = sum(f.stat().st_size for f in dst_path.rglob("*") if f.is_file())
                print(f"  📁 Collected {src_path.name}: {size:,} bytes")
            except Exception as e:
                print(f"  ⚠️  Could not copy {src_path}: {e}")

        # Write file manifest
        manifest = []
        for f in files_dir.rglob("*"):
            if f.is_file():
                rel = f.relative_to(files_dir)
                manifest.append({"path": str(rel), "size": f.stat().st_size})
        self.files_size = sum(m["size"] for m in manifest)
        (files_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))

        return files_dir

    def _write_metadata(self, work_dir: Path) -> None:
        """Write backup metadata JSON."""
        metadata = {
            "backup_version": "1.0",
            "created_at": now_iso(),
            "hostname": os.uname().nodename,
            "target": self.target_name,
            "bucket": self.bucket,
            "components": {
                "database": bool(DB_URL),
                "files": bool(FILES_DIR),
            },
            "config": {
                "db_url_prefix": DB_URL.split("://")[0] if DB_URL else None,
                "files_dir": FILES_DIR or None,
                "retention_days": RETENTION_DAYS,
            },
        }
        (work_dir / "metadata.json").write_text(json.dumps(metadata, indent=2))

    def _archive(self, work_dir: Path) -> tuple[Path, int, str]:
        """Create compressed tar archive and compute checksum."""
        archive_path = BACKUP_DIR / BACKUP_FILENAME

        with tarfile.open(archive_path, "w:gz") as tar:
            tar.add(work_dir, arcname=self.target_name)

        # Checksum
        sha256 = hashlib.sha256()
        with open(archive_path, "rb") as f:
            while chunk := f.read(65536):
                sha256.update(chunk)
        checksum = sha256.hexdigest()

        # Write checksum file alongside archive
        checksum_path = BACKUP_DIR / CHECKSUM_FILENAME
        checksum_path.write_text(f"{checksum}  {BACKUP_FILENAME}\n")

        size = os.path.getsize(archive_path)
        print(f"  📦 Archive: {BACKUP_FILENAME} ({size:,} bytes, sha256: {checksum[:16]}...)")

        return archive_path, size, checksum

    def _upload(self, archive_path: Path) -> None:
        """Upload archive and checksum to S3 via aws-cli."""
        key = f"{self.prefix}{BACKUP_FILENAME}"
        checksum_key = f"{self.prefix}{CHECKSUM_FILENAME}"

        endpoint_flag = []
        if self.endpoint:
            endpoint_flag = ["--endpoint-url", self.endpoint]

        region_flag = ["--region", self.region]

        # Upload archive
        print(f"  ☁️  Uploading to s3://{self.bucket}/{key}")
        result = subprocess.run(
            ["aws", "s3", "cp", str(archive_path), f"s3://{self.bucket}/{key}"] + endpoint_flag + region_flag,
            capture_output=True, text=True, timeout=120,
        )
        if result.returncode != 0:
            raise RuntimeError(f"S3 upload failed: {result.stderr}")

        # Upload checksum
        checksum_path = BACKUP_DIR / CHECKSUM_FILENAME
        result = subprocess.run(
            ["aws", "s3", "cp", str(checksum_path), f"s3://{self.bucket}/{checksum_key}"] + endpoint_flag + region_flag,
            capture_output=True, text=True, timeout=30,
        )
        if result.returncode != 0:
            raise RuntimeError(f"S3 checksum upload failed: {result.stderr}")

        print(f"  ✅ Uploaded to s3://{self.bucket}/{key}")

    def _cleanup_old(self) -> None:
        """Remove local temporary files."""
        for f in BACKUP_DIR.glob("*.tar.gz"):
            if f.name != BACKUP_FILENAME:
                f.unlink()
        for f in BACKUP_DIR.glob("*.sha256"):
            if f.name != CHECKSUM_FILENAME:
                f.unlink()

        # Remote cleanup via S3 lifecycle is preferred, but we can also
        # list and delete old backups here if needed.
        print(f"  🧹 Retention: {RETENTION_DAYS} days")


def restore_backup(archive_name: str) -> None:
    """Restore portal from a backup archive stored in S3."""
    print(f"🔄 Restoring from {archive_name}")

    # Try primary bucket first, then DR
    for target_name, target in BACKUP_TARGETS.items():
        key = f"{target['prefix']}{archive_name}"
        endpoint_flag = ["--endpoint-url", target["endpoint"]] if target["endpoint"] else []

        print(f"  Looking in s3://{target['bucket']}/{key}...")
        result = subprocess.run(
            ["aws", "s3", "ls", f"s3://{target['bucket']}/{key}"] + endpoint_flag,
            capture_output=True, text=True, timeout=15,
        )
        if result.returncode == 0:
            # Download
            local_path = BACKUP_DIR / archive_name
            result = subprocess.run(
                ["aws", "s3", "cp", f"s3://{target['bucket']}/{key}", str(local_path)] + endpoint_flag,
                capture_output=True, text=True, timeout=120,
            )
            if result.returncode != 0:
                print(f"  ❌ Download failed: {result.stderr}")
                continue

            # Verify checksum
            checksum_key = f"{target['prefix']}{archive_name}.sha256"
            checksum_local = BACKUP_DIR / f"{archive_name}.sha256"
            result = subprocess.run(
                ["aws", "s3", "cp", f"s3://{target['bucket']}/{checksum_key}", str(checksum_local)] + endpoint_flag,
                capture_output=True, text=True, timeout=15,
            )
            if result.returncode == 0:
                import hashlib
                sha256 = hashlib.sha256()
                with open(local_path, "rb") as f:
                    while chunk := f.read(65536):
                        sha256.update(chunk)
                expected = checksum_local.read_text().strip().split()[0]
                if sha256.hexdigest() != expected:
                    print(f"  ⚠️  Checksum mismatch! Expected {expected}, got {sha256.hexdigest()}")

            # Extract
            print(f"  📂 Extracting {local_path}...")
            with tarfile.open(local_path, "r:gz") as tar:
                tar.extractall(BACKUP_DIR)

            print(f"\n✅ Restored to {BACKUP_DIR}/{target_name}/")
            print(f"   To restore database:")
            if DB_URL:
                print(f"   cat {BACKUP_DIR}/{target_name}/portal-database.sql.gz | gunzip | sqlite3 portal.db")
            print(f"   To restore files:")
            print(f"   cp -r {BACKUP_DIR}/{target_name}/files/* /path/to/portal/")
            return

    print(f"❌ Could not find {archive_name} in any backup target")


# ─── Main ──────────────────────────────────────────────────────────

def main():
    import argparse

    parser = argparse.ArgumentParser(description="IT Pro Partner Portal Backup")
    parser.add_argument("--target", choices=["primary", "cross_region", "all"], default="all",
                        help="Backup target (default: all configured)")
    parser.add_argument("--restore", metavar="ARCHIVE.tar.gz",
                        help="Restore from a backup archive")
    parser.add_argument("--list", action="store_true",
                        help="List available backups in S3")
    args = parser.parse_args()

    if not BACKUP_TARGETS:
        print("❌ No backup targets configured. Set S3_BUCKET_PRIMARY or S3_BUCKET_DR.")
        sys.exit(1)

    if args.restore:
        restore_backup(args.restore)
        return

    if args.list:
        for target_name, target in BACKUP_TARGETS.items():
            endpoint_flag = ["--endpoint-url", target["endpoint"]] if target["endpoint"] else []
            print(f"📋 Backups in {target['bucket']} ({target_name}):")
            result = subprocess.run(
                ["aws", "s3", "ls", f"s3://{target['bucket']}/{target['prefix']}"] + endpoint_flag,
                capture_output=True, text=True, timeout=15,
            )
            if result.returncode == 0:
                for line in result.stdout.strip().split("\n"):
                    if line:
                        print(f"  {line}")
            else:
                print(f"  {result.stderr}")
        return

    # Run backups
    results = []
    for target_name, target_config in BACKUP_TARGETS.items():
        if args.target != "all" and args.target != target_name:
            continue

        print(f"\n{'='*60}")
        print(f"Backing up to {target_name}: {target_config['bucket']}")
        print(f"{'='*60}")

        runner = BackupRunner(target_name, target_config)
        try:
            result = runner.run()
            results.append(result)
            print(f"  ✅ Success: {result['size_bytes']:,} bytes in {result['duration_seconds']:.1f}s")
        except Exception as e:
            print(f"  ❌ Failed: {e}")
            results.append({"target": target_name, "status": "failed", "error": str(e)})

    # Print summary
    print(f"\n{'='*60}")
    print("Backup Summary")
    print(f"{'='*60}")
    for r in results:
        status_icon = "✅" if r["status"] == "success" else "❌"
        print(f"  {status_icon} {r['target']}: {r['status']}")
        if r["status"] == "success":
            print(f"     Archive: {r['archive_key']}")
            print(f"     Size: {r['size_bytes']:,} bytes")
            print(f"     Duration: {r['duration_seconds']:.1f}s")
        else:
            print(f"     Error: {r.get('error', 'unknown')}")

    if DRY_RUN:
        print(f"\n🔍 DRY RUN — No backups were uploaded")


if __name__ == "__main__":
    # Check for DRY_RUN
    DRY_RUN = os.environ.get("DRY_RUN", "").lower() in ("1", "true", "yes")
    main()