import hashlib
import importlib.util
import os
from pathlib import Path
import sqlite3
import struct
import tempfile
import unittest
from unittest import mock


PLUGIN_ROOT = Path(__file__).resolve().parents[1]
SCRIPT_PATH = (
    PLUGIN_ROOT
    / "skills"
    / "rodgers-ios-music-export"
    / "scripts"
    / "ios_music_export.py"
)
SPEC = importlib.util.spec_from_file_location("ios_music_export", SCRIPT_PATH)
MODULE = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(MODULE)


def create_media_database(path):
    connection = sqlite3.connect(path)
    connection.executescript(
        """
        CREATE TABLE item (
            item_pid INTEGER PRIMARY KEY,
            genre_id INTEGER,
            item_artist_pid INTEGER,
            album_pid INTEGER,
            base_location_id INTEGER
        );
        CREATE TABLE item_extra (
            item_pid INTEGER PRIMARY KEY,
            title TEXT,
            location TEXT,
            file_size INTEGER
        );
        CREATE TABLE genre (genre_id INTEGER PRIMARY KEY, genre TEXT);
        CREATE TABLE item_artist (item_artist_pid INTEGER PRIMARY KEY, item_artist TEXT);
        CREATE TABLE album (album_pid INTEGER PRIMARY KEY, album TEXT);
        CREATE TABLE base_location (base_location_id INTEGER PRIMARY KEY, path TEXT);
        """
    )
    connection.executemany(
        "INSERT INTO genre VALUES (?, ?)",
        [(1, "Electro"), (2, "electro"), (3, "Elektronisch"), (4, "Electro House")],
    )
    connection.execute("INSERT INTO item_artist VALUES (10, 'Test/Artist')")
    connection.execute("INSERT INTO album VALUES (20, 'Test Album')")
    connection.execute("INSERT INTO base_location VALUES (30, 'iTunes_Control/Music')")
    rows = [
        (101, 1, "Same: Track", "F00/AAAA.m4a", 5),
        (102, 2, "Same: Track", "F01/BBBB.M4A", 6),
        (103, 3, "Andere track", "F02/CCCC.m4a", 7),
        (104, 4, "House track", "F03/DDDD.m4a", 8),
    ]
    for item_pid, genre_id, title, location, size in rows:
        connection.execute("INSERT INTO item VALUES (?, ?, 10, 20, 30)", (item_pid, genre_id))
        connection.execute(
            "INSERT INTO item_extra VALUES (?, ?, ?, ?)",
            (item_pid, title, location, size),
        )
    connection.commit()
    connection.close()


class FakeAFCClient:
    def __init__(self, content):
        self.content = content
        self.download_calls = []

    def download(self, remote_path, local_path, destination_fd=None):
        self.download_calls.append(remote_path)
        data = self.content[remote_path]
        flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
        if hasattr(os, "O_NOFOLLOW"):
            flags |= os.O_NOFOLLOW
        descriptor = os.open(local_path, flags, 0o600, dir_fd=destination_fd)
        with os.fdopen(descriptor, "wb") as output:
            output.write(data)
        return len(data), hashlib.sha256(data).hexdigest()

    def get_file_info(self, remote_path):
        if remote_path not in self.content:
            raise MODULE.AFCError(MODULE.AFC_E_OBJECT_NOT_FOUND, MODULE.AFC_OP_GET_FILE_INFO)
        data = self.content[remote_path]
        return {
            "st_ifmt": "S_IFREG",
            "st_size": str(len(data)),
            "st_mtime": "1",
            "st_blocks": "1",
        }


def sample_plan():
    return [
        {
            "item_pid": "1",
            "genre": "Electro",
            "source_device_fingerprint": "a" * 64,
            "database_snapshot_sha256": "b" * 64,
            "title": "One",
            "artist": "Artist",
            "album": "Album",
            "file_size": 5,
            "source_mtime": "1",
            "source_blocks": "1",
            "remote_path": "/iTunes_Control/Music/F00/one.m4a",
            "filename": "Artist — One.m4a",
        },
        {
            "item_pid": "2",
            "genre": "Electro",
            "source_device_fingerprint": "a" * 64,
            "database_snapshot_sha256": "b" * 64,
            "title": "Two",
            "artist": "Artist",
            "album": "Album",
            "file_size": 6,
            "source_mtime": "1",
            "source_blocks": "1",
            "remote_path": "/iTunes_Control/Music/F01/two.mp3",
            "filename": "Artist — Two.mp3",
        },
    ]


class IOSMusicExportTests(unittest.TestCase):
    def test_exact_case_insensitive_genre_and_collision_names(self):
        with tempfile.TemporaryDirectory() as directory:
            database = Path(directory) / "MediaLibrary.sqlitedb"
            create_media_database(database)
            plan = MODULE.make_export_plan(database, "ELECTRO")

        self.assertEqual([row["item_pid"] for row in plan], ["101", "102"])
        self.assertEqual(plan[0]["filename"], "Test⁄Artist — Same - Track.m4a")
        self.assertEqual(plan[1]["filename"], "Test⁄Artist — Same - Track [2].m4a")
        self.assertTrue(plan[0]["remote_path"].endswith("/F00/AAAA.m4a"))

    def test_adjacent_genres_are_not_selected(self):
        with tempfile.TemporaryDirectory() as directory:
            database = Path(directory) / "MediaLibrary.sqlitedb"
            create_media_database(database)
            plan = MODULE.make_export_plan(database, "Elektronisch")

        self.assertEqual(len(plan), 1)
        self.assertEqual(plan[0]["item_pid"], "103")

    def test_unknown_schema_fails_closed(self):
        with tempfile.TemporaryDirectory() as directory:
            database = Path(directory) / "MediaLibrary.sqlitedb"
            sqlite3.connect(database).close()
            with self.assertRaisesRegex(RuntimeError, "unsupported MediaLibrary schema"):
                MODULE.make_export_plan(database, "Electro")

    def test_two_identical_database_captures_are_consolidated_and_checked(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = root / "source.sqlitedb"
            create_media_database(source)
            client = FakeAFCClient({MODULE.MEDIA_DATABASE_REMOTE: source.read_bytes()})
            snapshot = MODULE.snapshot_media_database(client, root / "snapshot-work")
            plan = MODULE.make_export_plan(snapshot, "Electro")

        self.assertEqual(len(plan), 2)
        self.assertEqual(client.download_calls.count(MODULE.MEDIA_DATABASE_REMOTE), 2)

    def test_export_hashes_verifies_and_resumes(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = root / "manifest.tsv"
            plan = sample_plan()
            client = FakeAFCClient({
                "/iTunes_Control/Music/F00/one.m4a": b"12345",
                "/iTunes_Control/Music/F01/two.mp3": b"abcdef",
            })
            first = MODULE.export_plan(client, plan, destination, manifest)
            verification = MODULE.verify_export(
                manifest, destination, "Electro", probe_audio=False
            )
            second_client = FakeAFCClient({})
            second = MODULE.export_plan(second_client, plan, destination, manifest)

        self.assertEqual(first["copied_tracks"], 2)
        self.assertEqual(verification["manifest_tracks"], 2)
        self.assertEqual(verification["hash_errors"], [])
        self.assertEqual(second["resumed_tracks"], 2)
        self.assertEqual(second_client.download_calls, [])

    def test_untrusted_existing_file_is_preserved(self):
        with tempfile.TemporaryDirectory() as directory:
            destination = Path(directory).resolve() / "music"
            destination.mkdir()
            (destination / "Artist — One.m4a").write_bytes(b"12345")
            plan = [sample_plan()[0]]
            with self.assertRaisesRegex(FileExistsError, "preserving untrusted"):
                MODULE.inspect_existing_destination(destination, plan, {})

    def test_verify_requires_present_nonempty_manifest_and_correct_genre(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            destination.mkdir()
            manifest = root / "manifest.tsv"
            with self.assertRaisesRegex(RuntimeError, "manifest does not exist"):
                MODULE.verify_export(manifest, destination, "Electro", probe_audio=False)
            MODULE.write_manifest_atomic(manifest, [])
            with self.assertRaisesRegex(RuntimeError, "manifest is empty"):
                MODULE.verify_export(manifest, destination, "Electro", probe_audio=False)

            client = FakeAFCClient({
                "/iTunes_Control/Music/F00/one.m4a": b"12345",
                "/iTunes_Control/Music/F01/two.mp3": b"abcdef",
            })
            MODULE.export_plan(client, sample_plan(), destination, manifest)
            with self.assertRaisesRegex(RuntimeError, "wrong_genre"):
                MODULE.verify_export(manifest, destination, "Jazz", probe_audio=False)

    def test_pending_manifest_never_trusts_same_size_final_file(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            destination.mkdir()
            manifest = root / "manifest.tsv"
            row = sample_plan()[0]
            (destination / row["filename"]).write_bytes(b"WRONG")
            pending = {
                **row,
                "partial_name": "",
                "actual_size": "0",
                "sha256": "",
                "status": "pending",
            }
            MODULE.write_manifest_atomic(manifest, [pending])
            with self.assertRaisesRegex(RuntimeError, "no trusted source hash"):
                MODULE.inspect_existing_destination(destination, [row], MODULE.load_manifest(manifest))

            pending["sha256"] = hashlib.sha256(b"WRONG").hexdigest()
            MODULE.write_manifest_atomic(manifest, [pending])
            with self.assertRaisesRegex(RuntimeError, "status is not complete"):
                MODULE.inspect_existing_destination(destination, [row], MODULE.load_manifest(manifest))

    def test_late_collision_is_not_overwritten(self):
        class RacingClient(FakeAFCClient):
            def __init__(self, content, final_path):
                super().__init__(content)
                self.final_path = final_path

            def download(self, remote_path, local_path, destination_fd=None):
                self.final_path.write_bytes(b"RIVAL")
                return super().download(
                    remote_path, local_path, destination_fd=destination_fd
                )

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = root / "manifest.tsv"
            row = sample_plan()[0]
            final_path = destination / row["filename"]
            client = RacingClient({row["remote_path"]: b"12345"}, final_path)
            with self.assertRaisesRegex(RuntimeError, "refusing to replace"):
                MODULE.export_plan(client, [row], destination, manifest)
            self.assertEqual(final_path.read_bytes(), b"RIVAL")

    def test_destination_symlink_swap_cannot_redirect_track_write(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            moved_destination = root / "music-original"
            escape = root / "escape"
            escape.mkdir()
            manifest = root / "manifest.tsv"
            row = sample_plan()[0]

            class SwappingClient(FakeAFCClient):
                def download(self, remote_path, local_path, destination_fd=None):
                    destination.rename(moved_destination)
                    destination.symlink_to(escape, target_is_directory=True)
                    return super().download(
                        remote_path, local_path, destination_fd=destination_fd
                    )

            client = SwappingClient({row["remote_path"]: b"12345"})
            with self.assertRaisesRegex(RuntimeError, "destination path changed concurrently"):
                MODULE.export_plan(client, [row], destination, manifest)
            self.assertFalse((escape / row["filename"]).exists())
            self.assertEqual((moved_destination / row["filename"]).read_bytes(), b"12345")
            self.assertEqual(MODULE.load_manifest(manifest)[row["filename"]]["status"], "pending")

    def test_download_failure_preserves_later_resume_records(self):
        class FailingClient(FakeAFCClient):
            def download(self, remote_path, local_path, destination_fd=None):
                raise IOError("injected download failure")

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = root / "manifest.tsv"
            plan = sample_plan()
            content = {
                "/iTunes_Control/Music/F00/one.m4a": b"12345",
                "/iTunes_Control/Music/F01/two.mp3": b"abcdef",
            }
            MODULE.export_plan(FakeAFCClient(content), plan, destination, manifest)
            (destination / plan[0]["filename"]).unlink()

            with self.assertRaisesRegex(IOError, "injected"):
                MODULE.export_plan(FailingClient(content), plan, destination, manifest)
            records = MODULE.load_manifest(manifest)
            self.assertEqual(set(records), {plan[0]["filename"], plan[1]["filename"]})
            self.assertTrue(records[plan[1]["filename"]]["sha256"])

            summary = MODULE.export_plan(FakeAFCClient(content), plan, destination, manifest)
            self.assertEqual(summary["copied_tracks"], 1)
            self.assertEqual(summary["resumed_tracks"], 1)

    def test_interrupted_direct_final_is_preserved_and_not_trusted(self):
        class PartialFailureClient(FakeAFCClient):
            def download(self, remote_path, local_path, destination_fd=None):
                flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
                if hasattr(os, "O_NOFOLLOW"):
                    flags |= os.O_NOFOLLOW
                descriptor = os.open(
                    local_path, flags, 0o600, dir_fd=destination_fd
                )
                with os.fdopen(descriptor, "wb") as output:
                    output.write(b"12")
                raise IOError("injected mid-track interruption")

        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = root / "manifest.tsv"
            row = sample_plan()[0]
            client = PartialFailureClient({row["remote_path"]: b"12345"})
            with self.assertRaisesRegex(IOError, "mid-track interruption"):
                MODULE.export_plan(client, [row], destination, manifest)
            final_path = destination / row["filename"]
            self.assertEqual(final_path.read_bytes(), b"12")
            with self.assertRaisesRegex(RuntimeError, "no trusted source hash"):
                MODULE.export_plan(FakeAFCClient({row["remote_path"]: b"12345"}), [row], destination, manifest)
            self.assertEqual(final_path.read_bytes(), b"12")

    def test_ready_final_is_hash_verified_and_recovered(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            destination.mkdir()
            manifest = root / "manifest.tsv"
            row = sample_plan()[0]
            final_path = destination / row["filename"]
            final_path.write_bytes(b"12345")
            ready = {
                **row,
                "partial_name": "",
                "actual_size": "5",
                "sha256": hashlib.sha256(b"12345").hexdigest(),
                "status": "ready_to_commit",
            }
            MODULE.write_manifest_atomic(manifest, [ready])

            summary = MODULE.export_plan(FakeAFCClient({}), [row], destination, manifest)
            verification = MODULE.verify_export(
                manifest, destination, "Electro", probe_audio=False
            )
            self.assertEqual(summary["recovered_tracks"], 1)
            self.assertEqual(verification["hash_errors"], [])
            self.assertEqual(final_path.read_bytes(), b"12345")

    def test_unrelated_manifest_is_preserved(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = root / "manifest.tsv"
            unrelated = {
                **sample_plan()[0],
                "filename": "Unrelated.m4a",
                "partial_name": "",
                "actual_size": "5",
                "sha256": "c" * 64,
                "status": "copied",
            }
            MODULE.write_manifest_atomic(manifest, [unrelated])
            original = manifest.read_bytes()
            with self.assertRaisesRegex(RuntimeError, "another export plan"):
                MODULE.export_plan(
                    FakeAFCClient({sample_plan()[0]["remote_path"]: b"12345"}),
                    [sample_plan()[0]],
                    destination,
                    manifest,
                )
            self.assertEqual(manifest.read_bytes(), original)
            self.assertFalse(destination.exists())

    def test_late_unrelated_manifest_collision_is_preserved(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = root / "manifest.tsv"
            row = sample_plan()[0]
            unrelated = {
                **row,
                "filename": "Late unrelated.m4a",
                "partial_name": "",
                "actual_size": "5",
                "sha256": "c" * 64,
                "status": "copied",
            }
            original_inspect = MODULE.inspect_existing_destination
            collided_bytes = []

            def inject_manifest_collision(
                target, plan, previous_records, destination_descriptor=None
            ):
                result = original_inspect(
                    target,
                    plan,
                    previous_records,
                    destination_descriptor=destination_descriptor,
                )
                MODULE.write_manifest_atomic(manifest, [unrelated])
                with manifest.open("ab") as output:
                    output.write(b"torn-unrelated-tail")
                collided_bytes.append(manifest.read_bytes())
                return result

            with mock.patch.object(
                MODULE, "inspect_existing_destination", side_effect=inject_manifest_collision
            ):
                with self.assertRaisesRegex(RuntimeError, "another export plan"):
                    MODULE.export_plan(
                        FakeAFCClient({row["remote_path"]: b"12345"}),
                        [row],
                        destination,
                        manifest,
                    )
            self.assertEqual(manifest.read_bytes(), collided_bytes[0])
            with self.assertRaisesRegex(RuntimeError, "truncated final event"):
                MODULE.load_manifest(manifest)

    def test_fault_during_initial_manifest_write_installs_no_manifest(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            manifest = root / "manifest.tsv"
            row = {
                **sample_plan()[0],
                "partial_name": "",
                "actual_size": "0",
                "sha256": "",
                "status": "pending",
            }
            real_write = os.write

            def fail_before_first_byte(descriptor, payload):
                raise OSError(28, "injected no space")

            with mock.patch.object(MODULE.os, "write", side_effect=fail_before_first_byte):
                with self.assertRaisesRegex(OSError, "injected no space"):
                    MODULE.write_manifest_atomic(manifest, [row], expected_plan=[row])
            self.assertFalse(manifest.exists())

            def fail_after_partial_header(descriptor, payload):
                real_write(descriptor, payload[:7])
                raise OSError(28, "injected partial header")

            with mock.patch.object(MODULE.os, "write", side_effect=fail_after_partial_header):
                with self.assertRaisesRegex(OSError, "injected partial header"):
                    MODULE.write_manifest_atomic(manifest, [row], expected_plan=[row])
            self.assertFalse(manifest.exists())
            self.assertEqual(list(root.glob("*.partial")), [])

            MODULE.write_manifest_atomic(manifest, [row], expected_plan=[row])
            self.assertEqual(MODULE.load_manifest(manifest)[row["filename"]]["status"], "pending")

    def test_torn_trailing_manifest_event_is_recovered_under_export_lock(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = root / "manifest.tsv"
            row = sample_plan()[0]
            row["title"] = "Line one\nLine two"
            MODULE.export_plan(
                FakeAFCClient({row["remote_path"]: b"12345"}),
                [row],
                destination,
                manifest,
            )
            with manifest.open("ab") as output:
                output.write(b"torn-final-event")
            with self.assertRaisesRegex(RuntimeError, "truncated final event"):
                MODULE.load_manifest(manifest)

            summary = MODULE.export_plan(FakeAFCClient({}), [row], destination, manifest)
            verification = MODULE.verify_export(
                manifest, destination, "Electro", probe_audio=False
            )
            self.assertEqual(summary["resumed_tracks"], 1)
            self.assertEqual(verification["hash_errors"], [])
            self.assertTrue(manifest.read_bytes().endswith(b"\n"))
            self.assertNotIn(b"Line one\nLine two", manifest.read_bytes())

    def test_symlink_destination_and_unexpected_directory_are_rejected(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            real_destination = root / "real"
            real_destination.mkdir()
            linked_destination = root / "linked"
            linked_destination.symlink_to(real_destination, target_is_directory=True)
            with self.assertRaisesRegex(RuntimeError, "symlink component"):
                MODULE.inspect_existing_destination(linked_destination, sample_plan(), {})

            (real_destination / "unexpected-directory").mkdir()
            with self.assertRaisesRegex(RuntimeError, "entries outside"):
                MODULE.inspect_existing_destination(real_destination, sample_plan(), {})

    def test_media_paths_are_confined_to_approved_roots(self):
        self.assertEqual(
            MODULE.make_safe_remote_media_path("iTunes_Control/Music/F00", "ABCD.m4a"),
            "/iTunes_Control/Music/F00/ABCD.m4a",
        )
        with self.assertRaisesRegex(RuntimeError, "unsafe media database path"):
            MODULE.make_safe_remote_media_path("iTunes_Control/Music/F00", "../secret")
        with self.assertRaisesRegex(RuntimeError, "outside approved roots"):
            MODULE.make_safe_remote_media_path("Safari", "history.db")

    def test_protocol_dispatchers_block_unknown_operations(self):
        with self.assertRaisesRegex(RuntimeError, "blocked AFC operation"):
            MODULE.AFCClient(None).request(0xFFFF)
        writable_open = struct.pack("<Q", 3) + b"/iTunes_Control/Music/file\0"
        with self.assertRaisesRegex(RuntimeError, "non-read-only AFC file-open"):
            MODULE.AFCClient(None).request(MODULE.AFC_OP_FILE_OPEN, writable_open)
        with self.assertRaisesRegex(RuntimeError, "blocked lockdownd request"):
            MODULE.lockdown_request(None, "Pair")
        with self.assertRaisesRegex(RuntimeError, "blocked non-read-only usbmuxd"):
            MODULE.usbmux_request({"MessageType": "SavePairRecord"})

    def test_lockdown_session_cannot_downgrade_paired_tls(self):
        class FakeSocket:
            def __init__(self):
                self.closed = False

            def close(self):
                self.closed = True

        fake_socket = FakeSocket()
        with (
            mock.patch.object(MODULE, "usbmux_connect", return_value=fake_socket),
            mock.patch.object(
                MODULE,
                "lockdown_request",
                return_value={"EnableSessionSSL": False, "SessionID": "session"},
            ),
            mock.patch.object(MODULE, "make_ssl_context") as make_ssl_context,
        ):
            with self.assertRaisesRegex(RuntimeError, "required paired TLS"):
                MODULE.start_lockdown_session({"DeviceID": 1}, {"HostID": "h", "SystemBUID": "s"})
        self.assertTrue(fake_socket.closed)
        make_ssl_context.assert_not_called()

    def test_hot_rollback_journal_is_rejected(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = root / "source.sqlitedb"
            create_media_database(source)
            client = FakeAFCClient(
                {
                    MODULE.MEDIA_DATABASE_REMOTE: source.read_bytes(),
                    MODULE.MEDIA_DATABASE_REMOTE + "-journal": b"hot rollback journal",
                }
            )
            snapshot = MODULE.capture_database_set(client, root / "snapshot-work")
        self.assertIsNone(snapshot)

    def test_invalid_nested_manifest_creates_no_destination_or_lock(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            destination = root / "music"
            manifest = destination / "manifest.tsv"
            with self.assertRaisesRegex(RuntimeError, "manifest outside"):
                MODULE.export_plan(FakeAFCClient({}), [sample_plan()[0]], destination, manifest)
            self.assertFalse(destination.exists())
            self.assertFalse(Path(str(manifest) + ".lock").exists())


if __name__ == "__main__":
    unittest.main()