#!/usr/bin/env python3

import os
import sys
import json
import sqlite3
import datetime
import shutil
import subprocess
import argparse
from collections import defaultdict

try:
    import boto3
    import lz4.frame
except ImportError:
    print("Error: Required Python packages 'boto3' or 'lz4' are not installed.")
    print("Please run: pip3 install boto3 lz4")
    sys.exit(1)

def main(args):
    """Main function to orchestrate the backup process."""
    print(f"Starting Open WebUI backup process at {datetime.datetime.now().isoformat()}")

    # 1. Setup temporary directory for staging backup files
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H%M%S")
    backup_staging_dir = os.path.join(args.tmp_dir, f"openwebui_backup_{timestamp}")
    db_local_path = os.path.join(backup_staging_dir, "webui.db")

    try:
        os.makedirs(backup_staging_dir, exist_ok=True)
        print(f"Created temporary staging directory: {backup_staging_dir}")

        # 2. Copy database from Docker container
        copy_db_from_docker(args.container_name, db_local_path)

        # 3. Export data from the database
        export_data_from_db(db_local_path, backup_staging_dir)

        # 4. Compress the backup directory
        archive_path = compress_backup(backup_staging_dir, timestamp)

        # 5. Upload to S3
        upload_to_s3(archive_path, args.s3_bucket, args.s3_region, args.s3_prefix)

    except Exception as e:
        print(f"An error occurred during the backup process: {e}", file=sys.stderr)
        sys.exit(1)
    finally:
        # 6. Cleanup local files
        print("Cleaning up local temporary files...")
        if os.path.exists(backup_staging_dir):
            shutil.rmtree(backup_staging_dir)
        if 'archive_path' in locals() and os.path.exists(archive_path):
            os.remove(archive_path)
        print("Cleanup complete.")

    print(f"Backup process finished successfully at {datetime.datetime.now().isoformat()}")

def copy_db_from_docker(container_name, local_path):
    """Copies the SQLite database from the Docker container to a local path."""
    db_container_path = "/app/backend/data/webui.db"
    command = ["docker", "cp", f"{container_name}:{db_container_path}", local_path]
    print(f"Executing: {' '.join(command)}")
    
    try:
        result = subprocess.run(command, check=True, capture_output=True, text=True)
        print("Successfully copied database from Docker container.")
    except FileNotFoundError:
        print("Error: 'docker' command not found. Is Docker installed and in your PATH?", file=sys.stderr)
        raise
    except subprocess.CalledProcessError as e:
        print(f"Error copying database from Docker container '{container_name}'.", file=sys.stderr)
        print(f"Stderr: {e.stderr}", file=sys.stderr)
        print("Please ensure the container is running and the name is correct.", file=sys.stderr)
        raise

def export_data_from_db(db_path, output_dir):
    """Connects to the SQLite DB and exports chats and config."""
    print("Connecting to the database and exporting data...")
    conn = None
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        # Export all user chats
        cursor.execute("""
            SELECT user.email, user.name, chat.title, chat.chat 
            FROM chat JOIN user ON chat.user_id = user.id
        """)
        rows = cursor.fetchall()

        user_chats = defaultdict(list)
        for email, name, title, chat_json_str in rows:
            try:
                chat_content = json.loads(chat_json_str)
            except json.JSONDecodeError:
                chat_content = {"error": "Failed to parse chat JSON."}
            
            user_chats[email].append({"title": title, "messages": chat_content})
        
        for email, chats in user_chats.items():
            safe_email = email.replace('@', '_at_').replace('.', '_')
            filename = os.path.join(output_dir, f"chats_{safe_email}.json")
            with open(filename, "w") as f:
                json.dump(chats, f, indent=2)
        print(f"Exported chats for {len(user_chats)} users.")

        # Export config
        cursor.execute("SELECT data FROM config ORDER BY id DESC LIMIT 1")
        config_row = cursor.fetchone()
        if config_row:
            config_data = json.loads(config_row[0])
            with open(os.path.join(output_dir, "openwebui_config.json"), "w") as f:
                json.dump(config_data, f, indent=2)
            print("Successfully exported configuration.")

    finally:
        if conn:
            conn.close()

def compress_backup(source_dir, timestamp):
    """Creates a .tar archive and then compresses it with lz4."""
    print("Compressing backup files...")
    base_name = f"openwebui_backup_{timestamp}"
    tar_path = os.path.join(os.path.dirname(source_dir), f"{base_name}.tar")
    lz4_path = f"{tar_path}.lz4"

    # Create a tar archive first
    shutil.make_archive(base_name=tar_path.replace('.tar',''), format='tar', root_dir=source_dir)
    
    # Compress the tar file with lz4
    with open(tar_path, 'rb') as f_in, open(lz4_path, 'wb') as f_out:
        f_out.write(lz4.frame.compress(f_in.read()))
    
    # Clean up the intermediate tar file
    os.remove(tar_path)
    print(f"Successfully created compressed archive: {lz4_path}")
    return lz4_path

def upload_to_s3(file_path, bucket, region, prefix):
    """Uploads a file to an S3 bucket."""
    s3_client = boto3.client("s3", region_name=region)
    s3_key = f"{prefix}{os.path.basename(file_path)}"
    
    print(f"Uploading {os.path.basename(file_path)} to s3://{bucket}/{s3_key}...")
    try:
        s3_client.upload_file(file_path, bucket, s3_key)
        print("Upload successful.")
    except Exception as e:
        print(f"Error uploading to S3: {e}", file=sys.stderr)
        raise

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Backup Open WebUI data to an S3 bucket.")
    parser.add_argument("--s3-bucket", required=True, help="S3 bucket name for backups.")
    parser.add_argument("--s3-region", required=True, help="AWS region of the S3 bucket.")
    parser.add_argument("--container-name", default="open-webui", help="Name of the Open WebUI Docker container.")
    parser.add_argument("--s3-prefix", default="openwebui_backups/", help="Prefix (folder) in the S3 bucket.")
    parser.add_argument("--tmp-dir", default="/tmp", help="Temporary directory for staging files.")
    
    args = parser.parse_args()
    main(args)