# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import os
import tempfile
import time

import mozfile
import mozlog
from marionette_harness import MarionetteTestCase


class BackupTestBase(MarionetteTestCase):
    """
    Base class for backup/recovery marionette tests.

    Provides common setup, teardown, and helper methods for testing
    backup and recovery scenarios involving selectable profiles.
    """

    _sandbox = "BackupTestSandbox"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.logger = mozlog.get_default_logger(component=self.__class__.__name__)

    def setUp(self):
        MarionetteTestCase.setUp(self)
        self.logger.info("Setting up test environment")

        # Profile backup is disabled while SQLite at-rest encryption is on
        # (per-database keys are not portable into the archive), so the backup
        # feature these tests exercise is unavailable then. We cannot just force
        # the pref off here: under a global encryption-on build the initial
        # session already created an encrypted profile, and restarting with
        # encryption off cannot reopen it. Skip instead -- on the shipping
        # (encryption-off) configuration this is a no-op and the tests run.
        if self.marionette.get_pref("security.storage.encryption.sqlite.enabled"):
            self.skipTest(
                "Profile backup is disabled when SQLite at-rest encryption is enabled"
            )

        self.marionette.enforce_gecko_prefs({
            "browser.backup.enabled": True,
            "browser.backup.log": True,
            "browser.backup.archive.enabled": True,
            "browser.backup.restore.enabled": True,
            "browser.backup.profiles.force-enable": True,
        })
        self.marionette.set_context("chrome")
        self.logger.info("Backup prefs configured")

        self._profile_name = None
        self._archive_path = None
        self._recovery_path = None
        self._new_profile_path = None
        self._new_profile_id = None
        self._cleanups = []

    def tearDown(self):
        self.logger.info("Tearing down test environment")
        self.marionette.restart(in_app=False, clean=True)

        for cleanup in self._cleanups:
            if cleanup.get("profile_name"):
                self.logger.info(f"Removing toolkit profile: {cleanup['profile_name']}")
                try:
                    self.marionette.execute_script(
                        """
                        let name = arguments[0];
                        let profileSvc = Cc["@mozilla.org/toolkit/profile-service;1"]
                            .getService(Ci.nsIToolkitProfileService);
                        let profile = profileSvc.getProfileByName(name);
                        profile.remove(false);
                        profileSvc.flush();
                        """,
                        script_args=(cleanup["profile_name"],),
                    )
                except Exception:
                    self.logger.warning(
                        f"Failed to remove profile: {cleanup['profile_name']}"
                    )

            if cleanup.get("path"):
                self.logger.info(f"Removing path: {cleanup['path']}")
                mozfile.remove(cleanup["path"])

        MarionetteTestCase.tearDown(self)
        self.logger.info("Teardown complete")

    def run_code(self, script, *args, **kwargs):
        """Run synchronous JS code."""
        return self.marionette.execute_script(
            script, new_sandbox=False, sandbox=self._sandbox, *args, **kwargs
        )

    def run_async_code(self, script, *args, **kwargs):
        """Run async JS code without error handling."""
        return self.marionette.execute_async_script(
            script, new_sandbox=False, sandbox=self._sandbox, *args, **kwargs
        )

    def run_async(self, script, script_args=None):
        """Run async JS code with error handling. Returns the script's return value."""
        wrapped = f"""
            let args = Array.from(arguments);
            let resolve = args.pop();
            (async () => {{
                try {{
                    return ["OK", await (async () => {{ {script} }})()];
                }} catch (e) {{
                    return ["ERROR", e.name, e.message, e.stack];
                }}
            }})().then(resolve);
        """
        result = self.marionette.execute_async_script(
            wrapped,
            script_args=script_args or [],
            new_sandbox=False,
            sandbox=self._sandbox,
        )
        self.assertEqual("OK", result[0], f"Script error: {result}")
        return result[1]

    def set_prefs(self, prefs):
        """Set prefs via Services.prefs (writes to profile's prefs.js only)."""
        self.marionette.execute_script(
            """
            for (let [name, value] of Object.entries(arguments[0])) {
                if (typeof value === "boolean")
                    Services.prefs.setBoolPref(name, value);
                else if (typeof value === "number")
                    Services.prefs.setIntPref(name, value);
                else if (typeof value === "string")
                    Services.prefs.setStringPref(name, value);
            }
            """,
            script_args=[prefs],
        )

    def register_profile_and_restart(self):
        """Register the profile with toolkit profile service and restart."""
        profile_name = "marionette-backup-test-" + str(int(time.time() * 1000))

        self.run_async(
            """
            let profD = Services.dirsvc.get("ProfD", Ci.nsIFile);
            let profileName = arguments[0];
            let profileSvc = Cc["@mozilla.org/toolkit/profile-service;1"]
                .getService(Ci.nsIToolkitProfileService);
            profileSvc.createProfile(profD, profileName, "backup-legacy-source");
            let { ProfileAge } = ChromeUtils.importESModule(
                "resource://gre/modules/ProfileAge.sys.mjs"
            );

            // Manually patch times.json as createNewProfile will not do this for existing directories.
            let profileAge = await ProfileAge(profD.path);
            profileAge._times.source = "backup-legacy-source";
            await profileAge.writeTimes();
            profileSvc.flush();
            """,
            script_args=[profile_name],
        )

        self.marionette.restart(clean=False, in_app=True)

        return profile_name

    def setup_selectable_profile(self):
        """Set up selectable profiles by creating a new profile in the database."""
        result = self.run_async(
            """
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            const { ProfileAge } = ChromeUtils.importESModule(
                "resource://gre/modules/ProfileAge.sys.mjs"
            );
            let newProfile = await SelectableProfileService.createNewProfile(false, null, "backup-selectable-source");
            let profileCount = (await SelectableProfileService.getAllProfiles()).length;
            let profile = SelectableProfileService.currentProfile;
            if (!profile) {
                throw new Error("currentProfile is null after createNewProfile");
            }

            // Manually patch times.json as createNewProfile will not do this for existing directories.
            let profileAge = await ProfileAge(PathUtils.profileDir);
            profileAge._times.source = "backup-selectable-source";
            await profileAge.writeTimes();
            return {
                path: profile.path,
                name: profile.name,
                store_id: SelectableProfileService.storeID,
                count: profileCount,
                id: profile.id
            };
            """
        )
        return result

    def create_backup(self):
        """Create a backup and return the archive path."""
        dest = os.path.join(tempfile.gettempdir(), "backup-dest")
        return self.run_async(
            """
            const { BackupService } = ChromeUtils.importESModule(
                "resource:///modules/backup/BackupService.sys.mjs"
            );
            let bs = BackupService.init();
            await bs.setParentDirPath(arguments[0]);
            let { archivePath } = await bs.createBackup();
            return archivePath;
            """,
            script_args=[dest],
        )

    def recover_backup(
        self, archive_path, recovery_path, replace_current_profile=False
    ):
        """Recover backup and return profile info."""
        return self.run_async(
            """
            const { BackupService } = ChromeUtils.importESModule(
                "resource:///modules/backup/BackupService.sys.mjs"
            );
            let [archivePath, recoveryPath, replaceCurrentProfile] = arguments;
            let bs = BackupService.get();
            let newProfileRoot = await IOUtils.createUniqueDirectory(
                PathUtils.tempDir, "backupTest"
            );
            let profile = await bs.recoverFromBackupArchive(
                archivePath, null, false, recoveryPath, newProfileRoot, replaceCurrentProfile
            );
            let rootDir = await profile.rootDir;
            return { name: profile.name, path: rootDir.path, id: profile.id };
            """,
            script_args=[archive_path, recovery_path, replace_current_profile],
        )

    def get_store_id(self):
        """Get the current storeID from SelectableProfileService."""
        return self.run_code(
            """
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            return SelectableProfileService.storeID;
            """
        )

    def get_store_id_pref(self):
        """Get the toolkit.profiles.storeID pref value."""
        return self.run_code(
            """
            return Services.prefs.getStringPref("toolkit.profiles.storeID", null);
            """
        )

    def has_selectable_profiles(self):
        """Check if selectable profiles have been created."""
        return self.run_code(
            """
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            return SelectableProfileService.hasCreatedSelectableProfiles();
            """
        )

    def wait_for_post_recovery(self):
        """Wait for post-recovery actions to complete."""
        self.run_async(
            """
            const { BackupService } = ChromeUtils.importESModule(
                "resource:///modules/backup/BackupService.sys.mjs"
            );
            await BackupService.get().postRecoveryComplete;
            """
        )

    def assert_profile_source(self, expected_source, profile_path=None):
        """Assert ProfileAge metadata for the current or specified profile."""
        source = self.run_async(
            """
            const { ProfileAge } = ChromeUtils.importESModule(
                "resource://gre/modules/ProfileAge.sys.mjs"
            );
            let profileAge = await ProfileAge(arguments[0]);
            return profileAge.source;
            """,
            script_args=[profile_path],
        )

        self.assertEqual(
            source,
            expected_source,
            f"ProfileAge source should be {expected_source}",
        )

    def get_all_profiles(self):
        """Get all profiles from the SelectableProfileService database."""
        return self.run_async(
            """
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            let profiles = await SelectableProfileService.getAllProfiles();
            return profiles.map(p => ({ id: p.id, name: p.name, path: p.path }));
            """
        )

    def get_original_profile_name(self):
        """Get the localized 'Original Profile' name."""
        return self.run_async(
            """
            let localization = new Localization(["browser/profiles.ftl"]);
            let [originalName] = await localization.formatMessages([
                { id: "original-profile-name" }
            ]);
            return originalName.value;
            """
        )

    def cleanup_selectable_profiles(self):
        """Clean up selectable profiles database."""
        self.run_async(
            """
            let [profileId] = arguments;
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            const { ProfilesDatastoreService } = ChromeUtils.importESModule(
                "moz-src:///toolkit/profile/ProfilesDatastoreService.sys.mjs"
            );

            if (profileId) {
                try {
                    let profile = await SelectableProfileService.getProfile(profileId);
                    if (profile) await SelectableProfileService.deleteProfile(profile);
                } catch (e) {}
            }

            let dbPath = await ProfilesDatastoreService.getProfilesStorePath();
            await SelectableProfileService.uninit();
            await ProfilesDatastoreService.uninit();
            for (let suffix of ["", "-shm", "-wal"]) {
                try { await IOUtils.remove(dbPath + suffix); } catch (e) {}
            }
            """,
            script_args=[self._new_profile_id],
        )

    def init_selectable_profile_service(self):
        """Initialize the SelectableProfileService."""
        self.run_async(
            """
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            await SelectableProfileService.init();
            """
        )

    def set_selectable_profile_metadata(self, name, avatar):
        """Set name and avatar on the current selectable profile."""
        self.run_async(
            """
            let [name, avatar] = arguments;
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            let profile = SelectableProfileService.currentProfile;
            if (!profile) {
                throw new Error("No current selectable profile");
            }
            profile.name = name;
            profile.avatar = avatar;
            """,
            script_args=[name, avatar],
        )

    def get_selectable_profile_metadata(self):
        """Return {name, avatar, theme} from the current selectable profile."""
        return self.run_async(
            """
            const { SelectableProfileService } = ChromeUtils.importESModule(
                "resource:///modules/profiles/SelectableProfileService.sys.mjs"
            );
            let profile = SelectableProfileService.currentProfile;
            if (!profile) {
                throw new Error("No current selectable profile");
            }
            return {
                name: profile.name,
                avatar: profile.avatar,
                theme: profile.theme,
            };
            """
        )