#!/usr/bin/env python3

# Libervia: an XMPP client
# Copyright (C) 2009-2026 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# The backup engine is asyncio based, but the test suite runs with the Twisted
# loop: the async flows are bridged with the "aio" decorator (as in
# memory/sqla.py).

import json
from pathlib import Path

from lxml import etree
from pytest import fixture, raises
from pytest_twisted import ensureDeferred as ed
from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from libervia.backend.memory.sqla_mapping import (
    Base,
    ParamGen,
    ParamInd,
    PrivateGen,
    PrivateGenBin,
    PrivateInd,
    Profile,
)
from libervia.backend.tools.backup import (
    FORMAT,
    FORMAT_VERSION,
    PIE_FILE_TPL,
    Backup,
    BackupError,
    ProfileData,
    ProfileHandler,
    ProfileMeta,
    Report,
    Staging,
    get_db_driver,
    sqlite_db_config,
)
from libervia.backend.tools.backup.core import (
    DB_FILE,
    NS_PARAMS,
    NS_PRIVATE,
    NS_PRIVATE_BIN,
)
from libervia.backend.tools.backup.xep_0227 import (
    NS_PUBSUB,
    NS_ROSTER,
    parse_pie,
    sections_by_ns,
)
from libervia.backend.tools.utils import aio

PEP_ITEM_DATA = (
    '<item id="item1"><nick xmlns="http://jabber.org/protocol/nick">Romy</nick></item>'
)
ROSTER_ITEM = (
    "<item xmlns='jabber:iq:roster' jid='contact@example.org' subscription='both'/>"
)


async def build_db(db_path: Path, data: bool = True) -> None:
    """Create a small database (one profile, optionally with data rows)"""
    db_path.unlink(missing_ok=True)
    engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    session_factory = async_sessionmaker(engine, expire_on_commit=False)
    async with session_factory() as session:
        session.add(Profile(id=1, name="test"))
        if data:
            session.add(
                ParamGen(category="Connection", name="Server", value="example.org")
            )
            session.add(
                ParamInd(
                    category="Connection",
                    name="JabberID",
                    profile_id=1,
                    value="user@example.org",
                )
            )
            session.add(PrivateGen(namespace="urn:xmpp:example", key="foo", value="bar"))
            session.add(
                PrivateInd(
                    namespace="urn:xmpp:example", key="baz", value="qux", profile_id=1
                )
            )
            session.add(
                PrivateGenBin(namespace="urn:xmpp:example", key="json", value={"a": 1})
            )
        await session.commit()
    await engine.dispose()


async def add_cache_data(db_path: Path) -> None:
    """Add a PEP node/item and a roster cache row to an existing database

    The pubsub rows are inserted with plain SQL (the app-specific column
    types don't need to be exercised here).
    """
    engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
    session_factory = async_sessionmaker(engine, expire_on_commit=False)
    try:
        async with engine.begin() as conn:
            await conn.execute(
                text(
                    "INSERT INTO pubsub_nodes (profile_id, service, name, pep, "
                    "subscribed, sync_state_updated, max_items) "
                    "VALUES (1, 'user@example.org', 'http://jabber.org/protocol/nick', "
                    "1, 0, 0.0, 0)"
                )
            )
            node_id = (
                await conn.execute(text("SELECT id FROM pubsub_nodes"))
            ).scalar_one()
            await conn.execute(
                text(
                    "INSERT INTO pubsub_items (node_id, name, data, created, updated) "
                    "VALUES (:node_id, 'item1', :data, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
                ),
                {"node_id": node_id, "data": PEP_ITEM_DATA},
            )
        async with session_factory() as session:
            session.add(
                PrivateInd(
                    namespace="urn:xmpp:features:rosterver",
                    key="@version@",
                    value="42",
                    profile_id=1,
                )
            )
            session.add(
                PrivateInd(
                    namespace="urn:xmpp:features:rosterver",
                    key="contact@example.org",
                    value=ROSTER_ITEM,
                    profile_id=1,
                )
            )
            await session.commit()
    finally:
        await engine.dispose()


def manifest_bytes(format=FORMAT, format_version=FORMAT_VERSION) -> bytes:
    """Get the bytes of a minimal valid manifest"""
    return json.dumps(
        {
            "format": format,
            "format_version": format_version,
            "app": "libervia",
            "app_version": "1.0",
            "date": "2026-01-01T00:00:00Z",
        }
    ).encode()


@fixture
def env(tmp_path):
    """Two local data directories: "src" and "fresh" (databases are built by the tests)"""
    src, fresh = tmp_path / "src", tmp_path / "fresh"
    for d in (src, fresh):
        d.mkdir()
        (d / "param").write_text("<params/>")
    return src, fresh


def profile_handlers(types: list[str]) -> list[ProfileHandler]:
    """Profile handlers of the given data types"""
    return [h for h in Backup.get_handlers(types) if isinstance(h, ProfileHandler)]


@fixture
def prof_handlers():
    """Profile handlers of the "params" and "private" data types"""
    return profile_handlers(["params", "private"])


class TestBackup:

    def test_profile_data(self):
        """The profile JID is parsed, and the profile name is the fallback"""
        data = ProfileData(name="test", id=1, jid="user@example.org")
        assert (data.host_jid, data.user_name) == ("example.org", "user")
        data = ProfileData(name="test", id=1, jid="not a jid")
        assert data.host_jid is None
        assert data.user_name == "test"

    @ed
    async def test_export_profile(self, env, prof_handlers):
        """Exporting a profile builds a pie document with its data sections"""
        src, _ = env

        @aio
        async def flow():
            await build_db(src / "libervia.db")
            ctx = await Backup.open_context(
                sqlite_db_config(src), src / "param", src, readonly=True
            )
            staging = Staging()
            try:
                profile = (await Backup.get_profiles(ctx))[0]
                assert profile.jid == "user@example.org"
                await Backup.export_profile(ctx, staging, profile, prof_handlers)
                host_jid, user, sections = parse_pie(
                    staging.path(PIE_FILE_TPL.format(profile.name)).read_bytes()
                )
                assert (host_jid, user) == ("example.org", "user")
                assert set(sections_by_ns(sections)) == {
                    NS_PARAMS,
                    NS_PRIVATE,
                    NS_PRIVATE_BIN,
                }
            finally:
                await ctx.close()
                staging.close()

        await flow()

    @ed
    async def test_import_profile(self, env, prof_handlers):
        """Importing a profile restores its data in another database"""
        src, fresh = env

        @aio
        async def flow():
            await build_db(src / "libervia.db")
            await build_db(fresh / "libervia.db", data=False)
            ctx = await Backup.open_context(
                sqlite_db_config(src), src / "param", src, readonly=True
            )
            staging = Staging()
            try:
                profile = (await Backup.get_profiles(ctx))[0]
                await Backup.export_profile(ctx, staging, profile, prof_handlers)
                await ctx.close()
                ctx = await Backup.open_context(
                    sqlite_db_config(fresh), fresh / "param", fresh
                )
                report = Report()
                await Backup.import_profile(
                    ctx, staging, profile.name, prof_handlers, report
                )
                async with ctx.session() as session:
                    server = (await session.execute(select(ParamGen.value))).scalar_one()
                    jid_val = (await session.execute(select(ParamInd.value))).scalar_one()
                    private = (
                        await session.execute(select(PrivateInd.value))
                    ).scalar_one()
                    private_bin = (
                        await session.execute(select(PrivateGenBin.value))
                    ).scalar_one()
                await ctx.close()
            finally:
                staging.close()
            assert server == "example.org"
            assert jid_val == "user@example.org"
            assert private == "qux"
            assert private_bin == {"a": 1}

        await flow()

    def test_archive_round_trip(self, tmp_path):
        """The archive can be read back with the same contents and entry modes"""
        staging = Staging()
        try:
            staging.add_bytes("server-data-test.xml", b"<pie/>")
            staging.add_bytes("param.xml", b"<params/>", mode=0o644)
            manifest = Backup.build_manifest(
                Backup.get_app_version(),
                None,
                {"test": ProfileMeta(jid="user@example.org", types=["params"])},
                staging,
            )
            arc = Backup.make_archive(staging, manifest, tmp_path)
        finally:
            staging.close()
        data = Staging()
        try:
            checked = Backup.read_archive(arc, data)
            assert checked.app == "libervia"
            assert data.path("server-data-test.xml").read_bytes() == b"<pie/>"
            assert data.path("server-data-test.xml").stat().st_mode & 0o777 == 0o600
            assert data.path("param.xml").stat().st_mode & 0o777 == 0o644
        finally:
            data.close()

    def test_tampered_archive(self, tmp_path):
        """An archive whose content doesn't match the manifest is rejected"""
        staging = Staging()
        try:
            staging.add_bytes("data.bin", b"original data")
            manifest = Backup.build_manifest(Backup.get_app_version(), None, {}, staging)
            staging.path("data.bin").write_bytes(b"tampered data")
            arc = Backup.make_archive(staging, manifest, tmp_path)
        finally:
            staging.close()
        data = Staging()
        try:
            with raises(BackupError, match="checksum mismatch"):
                Backup.read_archive(arc, data)
        finally:
            data.close()

    def test_manifest_validation(self):
        """The manifest is rejected if it is not from this backup format"""
        with raises(BackupError, match="unsupported format"):
            Backup._load_manifest(manifest_bytes(format="other"))
        with raises(BackupError, match="unsupported format version"):
            Backup._load_manifest(manifest_bytes(format_version=2))
        manifest = Backup._load_manifest(manifest_bytes())
        assert manifest.app == "libervia"

    @ed
    async def test_db_snapshot_restore(self, env):
        """The database restore replaces the database, with a pre-restore copy"""
        src, fresh = env

        @aio
        async def flow():
            await build_db(src / "libervia.db")
            await build_db(fresh / "libervia.db", data=False)
            snapshot = get_db_driver(sqlite_db_config(src))
            restore = get_db_driver(sqlite_db_config(fresh))
            staging = Staging()
            try:
                snapshot.snapshot(staging)
                assert DB_FILE in staging
                with raises(BackupError):
                    restore.restore(staging, force=False)
                restore.restore(staging, force=True)
                assert len(list(fresh.glob("libervia.db.pre-restore-*"))) == 1
                # the restored database has the source data (the fresh one had none)
                ctx = await Backup.open_context(
                    sqlite_db_config(fresh), fresh / "param", fresh, readonly=True
                )
                async with ctx.session() as session:
                    value = (await session.execute(select(PrivateGen.value))).scalar_one()
                await ctx.close()
                assert value == "bar"
            finally:
                staging.close()

        await flow()

    @ed
    async def test_pep_round_trip(self, env):
        """PEP items are exported to pie, and imported back into the pubsub cache"""
        src, fresh = env

        @aio
        async def flow():
            await build_db(src / "libervia.db")
            await add_cache_data(src / "libervia.db")
            await build_db(fresh / "libervia.db", data=False)
            handlers = profile_handlers(["pep"])
            staging = Staging()
            ctx = await Backup.open_context(
                sqlite_db_config(src), src / "param", src, readonly=True
            )
            try:
                profile = (await Backup.get_profiles(ctx))[0]
                await Backup.export_profile(ctx, staging, profile, handlers)
                await ctx.close()
                _, _, sections = parse_pie(
                    staging.path(PIE_FILE_TPL.format(profile.name)).read_bytes()
                )
                (section,) = sections_by_ns(sections)[NS_PUBSUB]
                (items,) = section.findall(f"{{{NS_PUBSUB}}}items")
                assert items.get("service") == "user@example.org"
                item = items.findall(f"{{{NS_PUBSUB}}}item")[0]
                assert item.get("id") == "item1"
                assert item.findtext("{http://jabber.org/protocol/nick}nick") == "Romy"
                report = Report()
                ctx = await Backup.open_context(
                    sqlite_db_config(fresh), fresh / "param", fresh
                )
                await Backup.import_profile(ctx, staging, profile.name, handlers, report)
                async with ctx.session() as session:
                    data = (
                        await session.execute(
                            text("SELECT data FROM pubsub_items WHERE name='item1'")
                        )
                    ).scalar_one()
                await ctx.close()
                if isinstance(data, str):
                    data = data.encode()
                stored = etree.fromstring(data)
                assert stored.findtext("{http://jabber.org/protocol/nick}nick") == "Romy"
            finally:
                staging.close()

        await flow()

    @ed
    async def test_roster_export(self, env):
        """The roster cache (rows in the database) is exported as a pie section"""
        src, _ = env

        @aio
        async def flow():
            await build_db(src / "libervia.db")
            await add_cache_data(src / "libervia.db")
            ctx = await Backup.open_context(
                sqlite_db_config(src), src / "param", src, readonly=True
            )
            staging = Staging()
            try:
                profile = (await Backup.get_profiles(ctx))[0]
                await Backup.export_profile(
                    ctx, staging, profile, profile_handlers(["roster"])
                )
                _, _, sections = parse_pie(
                    staging.path(PIE_FILE_TPL.format(profile.name)).read_bytes()
                )
                (section,) = sections_by_ns(sections)[NS_ROSTER]
                assert section.get("ver") == "42"
                item = section.find(f"{{{NS_ROSTER}}}item")
                assert item.get("jid") == "contact@example.org"
            finally:
                await ctx.close()
                staging.close()

        await flow()

    def test_verify_archive(self, tmp_path):
        """verify_archive reports the checksum status of each entry"""
        staging = Staging()
        try:
            staging.add_bytes("data.bin", b"original data")
            manifest = Backup.build_manifest(Backup.get_app_version(), None, {}, staging)
            arc = Backup.make_archive(staging, manifest, tmp_path)
            # the manifest hashes the original content: the tampered archive
            # must be reported as such
            staging.path("data.bin").write_bytes(b"tampered data")
            arc_t = Backup.make_archive(staging, manifest, tmp_path)
        finally:
            staging.close()
        _, results = Backup.verify_archive(arc)
        assert results["data.bin"]["checksum"] == "ok"
        _, results = Backup.verify_archive(arc_t)
        assert results["data.bin"]["checksum"] == "mismatch"