import ctypes
import datetime
import shutil
import struct
import warnings
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
from typing import Any, Callable, Optional, cast

from pymobiledevice3.exceptions import NotEnoughDiskSpaceError, PyMobileDevice3Exception
from pymobiledevice3.service_connection import ServiceConnection

SIZE_FORMAT = ">I"
CODE_FORMAT = ">B"
CODE_FILE_DATA = 0xC
CODE_ERROR_REMOTE = 0xB
CODE_ERROR_LOCAL = 0x6
CODE_SUCCESS = 0
FILE_TRANSFER_TERMINATOR = b"\x00\x00\x00\x00"
BULK_OPERATION_ERROR = -13
APPLE_EPOCH = 978307200
ERRNO_TO_DEVICE_ERROR = {
    2: -6,
    17: -7,
    20: -8,
    21: -9,
    62: -10,
    5: -11,
    28: -15,
}

DLMessage = Sequence[Any]
ProgressCallback = Callable[[Any], None]
DLHandler = Callable[[DLMessage], None]


class DeviceLink:
    def __init__(self, service: ServiceConnection, root_path: Path) -> None:
        self.service: ServiceConnection = service
        self.root_path: Path = root_path
        self._dl_handlers: dict[str, DLHandler] = {
            "DLMessageCreateDirectory": self.create_directory,
            "DLMessageUploadFiles": self.upload_files,
            "DLMessageGetFreeDiskSpace": self.get_free_disk_space,
            "DLMessageMoveItems": self.move_items,
            "DLMessageRemoveItems": self.remove_items,
            "DLMessageDownloadFiles": self.download_files,
            "DLContentsOfDirectory": self.contents_of_directory,
            "DLMessageCopyItem": self.copy_item,
            "DLMessagePurgeDiskSpace": self.purge_disk_space,
        }

    async def _sendall(self, payload: bytes) -> None:
        await self.service.sendall(payload)

    async def _recvall(self, size: int) -> bytes:
        return await self.service.recvall(size)

    async def dl_loop(self, progress_callback: Optional[ProgressCallback] = None) -> Any:
        def _noop(_: Any) -> None:
            return None

        callback: ProgressCallback = progress_callback if progress_callback is not None else _noop

        while True:
            message = await self.receive_message()
            command = message[0]

            if command in (
                "DLMessageDownloadFiles",
                "DLMessageMoveFiles",
                "DLMessageMoveItems",
                "DLMessageRemoveFiles",
                "DLMessageRemoveItems",
            ):
                callback(message[3])
            elif command == "DLMessageUploadFiles":
                callback(message[2])

            if command == "DLMessageProcessMessage":
                if not message[1]["ErrorCode"]:
                    return message[1].get("Content")
                else:
                    raise PyMobileDevice3Exception(f"Device link error: {message[1]}")
            await self._dl_handlers[command](message)

    async def version_exchange(self) -> None:
        dl_message_version_exchange = await self.receive_message()
        version_major = dl_message_version_exchange[1]
        await self.service.send_plist(["DLMessageVersionExchange", "DLVersionsOk", version_major])
        dl_message_device_ready = await self.receive_message()
        if dl_message_device_ready[0] != "DLMessageDeviceReady":
            raise PyMobileDevice3Exception("Device link didn't return ready state")

    async def send_process_message(self, message: Mapping[str, Any]) -> None:
        await self.service.send_plist(["DLMessageProcessMessage", message])

    async def download_files(self, message: DLMessage) -> None:
        status: dict[str, dict[str, Any]] = {}
        files = cast(Iterable[str], message[1])
        for file in files:
            await self._sendall(struct.pack(SIZE_FORMAT, len(file)))
            await self._sendall(file.encode())

            try:
                file_path = self.root_path / file

                # split into chunks, otherwise we may crash BackupAgent2 by OOM
                # https://github.com/doronz88/pymobiledevice3/issues/1165#issuecomment-2376815692
                chunk_size = 128 * 1024 * 1024  # 128 MB

                with file_path.open("rb") as file_handle:
                    while True:
                        chunk_data = file_handle.read(chunk_size)
                        if not chunk_data:
                            break
                        await self._sendall(struct.pack(SIZE_FORMAT, len(chunk_data) + struct.calcsize(CODE_FORMAT)))
                        await self._sendall(struct.pack(CODE_FORMAT, CODE_FILE_DATA) + chunk_data)

                buffer = struct.pack(SIZE_FORMAT, struct.calcsize(CODE_FORMAT)) + struct.pack(CODE_FORMAT, CODE_SUCCESS)
                await self._sendall(buffer)
            except OSError as e:
                status[file] = {
                    "DLFileErrorString": e.strerror,
                    "DLFileErrorCode": ctypes.c_uint64(ERRNO_TO_DEVICE_ERROR[e.errno]).value,
                }
                await self._sendall(struct.pack(SIZE_FORMAT, len(e.strerror) + struct.calcsize(CODE_FORMAT)))
                await self._sendall(struct.pack(CODE_FORMAT, CODE_ERROR_LOCAL) + e.strerror.encode())

        await self._sendall(FILE_TRANSFER_TERMINATOR)
        if status:
            await self.status_response(BULK_OPERATION_ERROR, "Multi status", status)
        else:
            await self.status_response(0)

    async def contents_of_directory(self, message: DLMessage) -> None:
        data = {}
        path = self.root_path / cast(str, message[1])
        for file in path.iterdir():
            ftype = "DLFileTypeUnknown"
            if file.is_dir():
                ftype = "DLFileTypeDirectory"
            if file.is_file():
                ftype = "DLFileTypeRegular"
            modifications_data = datetime.datetime.fromtimestamp(file.stat().st_mtime - APPLE_EPOCH)
            modifications_data = modifications_data.replace(tzinfo=None)
            data[file.name] = {
                "DLFileType": ftype,
                "DLFileSize": file.stat().st_size,
                "DLFileModificationDate": modifications_data,
            }
        await self.status_response(0, status_dict=data)

    async def upload_files(self, _message: DLMessage) -> None:
        while True:
            device_name = await self._prefixed_recv()
            if not device_name:
                break
            file_name = await self._prefixed_recv()
            (size,) = struct.unpack(SIZE_FORMAT, await self._recvall(struct.calcsize(SIZE_FORMAT)))
            (code,) = struct.unpack(CODE_FORMAT, await self._recvall(struct.calcsize(CODE_FORMAT)))
            size -= struct.calcsize(CODE_FORMAT)
            with open(self.root_path / file_name, "wb") as fd:
                while size and code == CODE_FILE_DATA:
                    fd.write(await self._recvall(size))
                    (size,) = struct.unpack(SIZE_FORMAT, await self._recvall(struct.calcsize(SIZE_FORMAT)))
                    (code,) = struct.unpack(CODE_FORMAT, await self._recvall(struct.calcsize(CODE_FORMAT)))
                    size -= struct.calcsize(CODE_FORMAT)
            if code == CODE_ERROR_REMOTE:
                # iOS 17 beta devices give this error for: backup_manifest.db
                error_message = (await self._recvall(size)).decode()
                warnings.warn(
                    f"Failed to fully upload: {file_name}. Device file name: {device_name}. Reason: {error_message}",
                    stacklevel=2,
                )
                continue
            assert code == CODE_SUCCESS
        await self.status_response(0)

    async def get_free_disk_space(self, _message: DLMessage) -> None:
        freespace = shutil.disk_usage(self.root_path).free
        await self.status_response(0, status_dict=freespace)

    async def move_items(self, message: DLMessage) -> None:
        items = cast(Mapping[str, str], message[1])
        for src, dst in items.items():
            dest = self.root_path / dst
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(self.root_path / src, dest)
        await self.status_response(0)

    async def copy_item(self, message: DLMessage) -> None:
        src = self.root_path / cast(str, message[1])
        dest = self.root_path / cast(str, message[2])
        dest.parent.mkdir(parents=True, exist_ok=True)
        if src.is_dir():
            shutil.copytree(src, dest)
        else:
            shutil.copy(src, dest)
        await self.status_response(0)

    async def purge_disk_space(self, _message: DLMessage) -> None:
        raise NotEnoughDiskSpaceError()

    async def remove_items(self, message: DLMessage) -> None:
        for path in cast(Iterable[str], message[1]):
            rm_path = self.root_path / path
            if rm_path.is_dir():
                shutil.rmtree(rm_path)
            else:
                rm_path.unlink(missing_ok=True)
        await self.status_response(0)

    async def create_directory(self, message: DLMessage) -> None:
        path = cast(str, message[1])
        (self.root_path / path).mkdir(parents=True, exist_ok=True)
        await self.status_response(0)

    async def status_response(self, status_code: int, status_str: str = "", status_dict: Any = None) -> None:
        await self.service.send_plist([
            "DLMessageStatusResponse",
            ctypes.c_uint64(status_code).value,
            status_str if status_str else "___EmptyParameterString___",
            status_dict if status_dict is not None else {},
        ])

    async def receive_message(self) -> DLMessage:
        return cast(DLMessage, await self.service.recv_plist())

    async def disconnect(self) -> None:
        await self.service.send_plist(["DLMessageDisconnect", "___EmptyParameterString___"])

    async def _prefixed_recv(self) -> str:
        (size,) = struct.unpack(SIZE_FORMAT, await self._recvall(struct.calcsize(SIZE_FORMAT)))
        return (await self._recvall(size)).decode()
