from asyncio import Lock, Queue, TaskGroup
from collections import deque
from collections.abc import AsyncIterator, Sequence
from contextlib import AsyncExitStack, asynccontextmanager
from dataclasses import dataclass, field, replace
from datetime import datetime, timezone
from itertools import tee
from pathlib import PurePath
from typing import override

from .exceptions import (
    AuthenticationError,
    InvalidServiceError,
    NodeExistsError,
    NodeNotFoundError,
)
from .lib import (
    else_none,
    is_valid_name,
    normalize_path,
)
from .types import (
    ChangeAction,
    CreateFileService,
    CreateHasher,
    CreateService,
    CreateServiceMiddleware,
    CreateSnapshotService,
    Drive,
    FileService,
    MediaInfo,
    Node,
    ReadableFile,
    Service,
    SnapshotService,
    SourceConfig,
    WritableFile,
)


_API_VERSION = 5

_VIRTUAL_ROOT_ID = ""


def compose_service[T: Service](
    base: CreateService[T],
    *middleware: CreateServiceMiddleware[T],
) -> CreateService[T]:
    @asynccontextmanager
    async def _factory() -> AsyncIterator[T]:
        async with AsyncExitStack() as stack:
            service = await stack.enter_async_context(base())
            if service.api_version != _API_VERSION:
                raise InvalidServiceError(
                    f"invalid version: required {_API_VERSION}, got {service.api_version}"
                )
            for mw in middleware:
                service = await stack.enter_async_context(mw(service))
                if service.api_version != _API_VERSION:
                    raise InvalidServiceError(
                        f"invalid version: required {_API_VERSION}, got {service.api_version}"
                    )
            yield service

    return _factory


@asynccontextmanager
async def create_multi_drive(
    sources: Sequence[SourceConfig],
) -> AsyncIterator[Drive]:
    """Create a configured Drive instance for cloud storage operations.

    Accepts a sequence of SourceConfig objects. With a single source,
    yields a single-mode Drive (paths unchanged). With multiple sources,
    yields a multi-mode Drive where each source's paths are prefixed with
    its name (e.g., /google/docs/file.txt).

    Args:
        sources: One or more SourceConfig objects describing storage backends.

    Yields:
        A fully configured Drive instance ready for operations.

    Raises:
        InvalidServiceError: If any service reports an incompatible API version.
        ValueError: If sources is empty.
    """
    if not sources:
        raise ValueError("at least one source is required")

    if len(sources) == 1:
        source = sources[0]
        async with (
            _create_service(source.snapshot) as snapshot_service,
            _create_service(source.file) as file_service,
        ):
            yield _SingleDrive(
                file_service=file_service, snapshot_service=snapshot_service
            )
    else:
        async with AsyncExitStack() as stack:
            source_states: dict[str, _SourceState] = {}
            for source in sources:
                ss = await stack.enter_async_context(_create_service(source.snapshot))
                fs = await stack.enter_async_context(_create_service(source.file))
                source_states[source.name] = _SourceState(
                    file_service=fs, snapshot_service=ss
                )
            yield _MultiDrive(source_states)


@asynccontextmanager
async def create_drive(
    *,
    file: CreateFileService,
    snapshot: CreateSnapshotService,
) -> AsyncIterator[Drive]:
    """Create a configured Drive instance for single-source usage.

    Args:
        file: Factory function returning a FileService context manager.
              Also accepts the output of compose_service().
        snapshot: Factory function returning a SnapshotService context manager.
                  Also accepts the output of compose_service().

    Yields:
        A fully configured Drive instance ready for operations.

    Raises:
        InvalidServiceError: If the service reports an incompatible API version.
    """
    async with create_multi_drive(
        sources=[SourceConfig(name="", file=file, snapshot=snapshot)]
    ) as drive:
        yield drive


@asynccontextmanager
async def _create_service[T: Service](create_service: CreateService[T]):
    async with create_service() as service:
        if service.api_version != _API_VERSION:
            raise InvalidServiceError(
                f"invalid version: required {_API_VERSION}, got {service.api_version}"
            )
        yield service


@dataclass
class _SourceState:
    file_service: FileService
    snapshot_service: SnapshotService
    lock: Lock = field(default_factory=Lock)


def _scope_id(source: str, original_id: str) -> str:
    return f"{source}:{original_id}"


def _parse_id(scoped_id: str) -> tuple[str, str]:
    source, _, original_id = scoped_id.partition(":")
    return source, original_id


def _scope_node(source: str, node: Node) -> Node:
    new_parent_id = (
        _VIRTUAL_ROOT_ID
        if node.parent_id is None
        else _scope_id(source, node.parent_id)
    )
    return replace(node, id=_scope_id(source, node.id), parent_id=new_parent_id)


def _unscope_node(source: str, node: Node) -> Node:
    _, original_id = _parse_id(node.id)
    if node.parent_id == _VIRTUAL_ROOT_ID:
        original_parent_id = None
    elif node.parent_id is None:
        original_parent_id = None
    else:
        _, original_parent_id = _parse_id(node.parent_id)
    return replace(node, id=original_id, parent_id=original_parent_id)


def _scope_change(source: str, change: ChangeAction) -> ChangeAction:
    match change:
        case (True, id_):
            return (True, _scope_id(source, id_))
        case (False, node):
            return (False, _scope_node(source, node))


def _make_virtual_root() -> Node:
    epoch = datetime(1970, 1, 1, tzinfo=timezone.utc)
    return Node(
        id=_VIRTUAL_ROOT_ID,
        parent_id=None,
        name="",
        is_directory=True,
        is_trashed=False,
        ctime=epoch,
        mtime=epoch,
        mime_type="",
        hash="",
        size=0,
        is_image=False,
        is_video=False,
        width=0,
        height=0,
        ms_duration=0,
        private=None,
    )


class _SingleDrive(Drive):
    def __init__(
        self,
        *,
        file_service: FileService,
        snapshot_service: SnapshotService,
    ) -> None:
        self._sync_lock = Lock()
        self._fs = file_service
        self._ss = snapshot_service

    @override
    async def get_root(self) -> Node:
        return await self._ss.get_root()

    @override
    async def get_node_by_id(self, node_id: str) -> Node:
        return await self._ss.get_node_by_id(node_id)

    @override
    async def get_node_by_path(self, path: PurePath) -> Node:
        path = normalize_path(path)
        return await self._ss.get_node_by_path(path)

    @override
    async def resolve_path(self, node: Node) -> PurePath:
        return await self._ss.resolve_path_by_id(node.id)

    @override
    async def get_child_by_name(
        self,
        name: str,
        parent: Node,
    ) -> Node:
        return await self._ss.get_child_by_name(name, parent.id)

    @override
    async def get_children(self, parent: Node) -> list[Node]:
        return await self._ss.get_children_by_id(parent.id)

    @override
    async def get_trashed_nodes(self, flatten: bool = False) -> list[Node]:
        rv = await self._ss.get_trashed_nodes()
        if flatten:
            return rv

        ancestor_set = set(_.id for _ in rv if _.is_directory)
        if not ancestor_set:
            return rv

        table = {_.id: _ for _ in rv}
        return [_ for _ in rv if not _in_ancestor_set(table, _, ancestor_set)]

    @override
    async def find_nodes_by_regex(self, pattern: str) -> list[Node]:
        return await self._ss.find_nodes_by_regex(pattern)

    @override
    async def walk(
        self,
        node: Node,
        *,
        include_trashed: bool = False,
    ) -> AsyncIterator[tuple[Node, list[Node], list[Node]]]:
        if not node.is_directory:
            return
        if node.is_trashed and not include_trashed:
            return

        q = deque([node])
        while q:
            node = q.popleft()
            children = await self.get_children(node)
            children = (_ for _ in children if not _.is_trashed or include_trashed)

            directories, files = tee(children, 2)
            directories = [_ for _ in directories if _.is_directory]
            files = [_ for _ in files if not _.is_directory]

            yield node, directories, files

            q.extend(directories)

    @override
    async def create_directory(
        self,
        name: str,
        parent: Node,
        *,
        exist_ok: bool = False,
    ) -> Node:
        # sanity check
        if not parent.is_directory:
            raise ValueError("parent is not a directory")
        if not name:
            raise ValueError("directory name is empty")
        if not is_valid_name(name):
            raise ValueError("no `/` or `\\` allowed in directory name")
        if not await self.is_authenticated():
            raise AuthenticationError()

        if not exist_ok:
            node = await else_none(
                self.get_child_by_name(
                    name,
                    parent,
                )
            )
            if node:
                raise NodeExistsError(node)

        return await self._fs.create_directory(
            name,
            parent,
            exist_ok=exist_ok,
            private=None,
        )

    @asynccontextmanager
    @override
    async def download_file(self, node: Node) -> AsyncIterator[ReadableFile]:
        # sanity check
        if node.is_directory:
            raise ValueError("node should be a file")
        if not await self.is_authenticated():
            raise AuthenticationError()

        async with self._fs.download_file(node) as file:
            yield file

    @asynccontextmanager
    @override
    async def upload_file(
        self,
        name: str,
        parent: Node,
        *,
        size: int | None = None,
        mime_type: str | None = None,
        media_info: MediaInfo | None = None,
    ) -> AsyncIterator[WritableFile]:
        # sanity check
        if not parent.is_directory:
            raise ValueError("parent is not a directory")
        if not name:
            raise ValueError("directory name is empty")
        if not is_valid_name(name):
            raise ValueError("no `/` or `\\` allowed in directory name")
        if not await self.is_authenticated():
            raise AuthenticationError()

        node = await else_none(self.get_child_by_name(name, parent))
        if node:
            raise NodeExistsError(node)

        async with self._fs.upload_file(
            name,
            parent,
            size=size,
            mime_type=mime_type,
            media_info=media_info,
            private=None,
        ) as file:
            yield file

    @override
    async def purge_trash(self) -> None:
        # sanity check
        if not await self.is_authenticated():
            raise AuthenticationError()

        await self._fs.purge_trash()

    @override
    async def delete(self, node: Node, *, permanent: bool = False) -> None:
        if not await self.is_authenticated():
            raise AuthenticationError()

        await self._fs.delete(node, permanent=permanent)

    @override
    async def restore(self, node: Node) -> Node:
        if not await self.is_authenticated():
            raise AuthenticationError()

        return await self._fs.restore(node)

    @override
    async def move(
        self,
        node: Node,
        *,
        new_parent: Node | None = None,
        new_name: str | None = None,
    ) -> Node:
        # sanity check
        if not await self.is_authenticated():
            raise AuthenticationError()

        if not new_parent and not new_name:
            raise ValueError("nothing to move")

        if new_name and not is_valid_name(new_name):
            raise ValueError("no `/` or `\\` allowed in file name")

        root_node = await self.get_root()
        if node.id == root_node.id:
            raise ValueError("root node is immutable")

        if new_parent:
            if not new_parent.is_directory:
                raise ValueError("new_parent is not a directory")
            if await _contains(self, node, new_parent):
                raise ValueError("new_parent is a descendant of the source node")

        return await self._fs.move(
            node,
            new_parent=new_parent,
            new_name=new_name,
        )

    @override
    async def sync(self) -> AsyncIterator[ChangeAction]:
        if not await self.is_authenticated():
            raise AuthenticationError()

        async with self._sync_lock:
            initial_cursor = await self._fs.get_initial_cursor()

            cursor = await self._ss.get_current_cursor()
            if not cursor:
                cursor = initial_cursor

            # no data before, get the root node and cache it
            if cursor == initial_cursor:
                node = await self._fs.get_root()
                await self._ss.set_root(node)

            async for changes, next_ in self._fs.get_changes(cursor):
                await self._ss.apply_changes(changes, next_)

                for change in changes:
                    yield change

    @override
    async def get_hasher_factory(self) -> CreateHasher:
        return await self._fs.get_hasher_factory()

    @override
    async def is_authenticated(self) -> bool:
        return await self._fs.is_authenticated()

    @override
    async def authenticate(self) -> None:
        return await self._fs.authenticate()


class _MultiDrive(Drive):
    def __init__(self, sources: dict[str, _SourceState]) -> None:
        self._sources = sources

    @override
    async def get_root(self) -> Node:
        return _make_virtual_root()

    @override
    async def get_node_by_id(self, node_id: str) -> Node:
        if node_id == _VIRTUAL_ROOT_ID:
            return _make_virtual_root()
        source_name, original_id = _parse_id(node_id)
        if source_name not in self._sources:
            raise NodeNotFoundError(node_id)
        state = self._sources[source_name]
        node = await state.snapshot_service.get_node_by_id(original_id)
        return _scope_node(source_name, node)

    @override
    async def get_node_by_path(self, path: PurePath) -> Node:
        path = normalize_path(path)
        parts = path.parts  # ('/', ...) or ('/',)

        if len(parts) == 1:
            return _make_virtual_root()

        source_name = parts[1]
        if source_name not in self._sources:
            raise NodeNotFoundError(str(path))
        state = self._sources[source_name]

        if len(parts) == 2:
            root = await state.snapshot_service.get_root()
            return replace(_scope_node(source_name, root), name=source_name)

        remaining = PurePath("/", *parts[2:])
        node = await state.snapshot_service.get_node_by_path(remaining)
        return _scope_node(source_name, node)

    @override
    async def resolve_path(self, node: Node) -> PurePath:
        if node.id == _VIRTUAL_ROOT_ID:
            return PurePath("/")
        source_name, original_id = _parse_id(node.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(node.id)
        if node.parent_id == _VIRTUAL_ROOT_ID:
            return PurePath("/", source_name)
        state = self._sources[source_name]
        path = await state.snapshot_service.resolve_path_by_id(original_id)
        return PurePath("/", source_name, *path.parts[1:])

    @override
    async def get_child_by_name(self, name: str, parent: Node) -> Node:
        if parent.id == _VIRTUAL_ROOT_ID:
            if name not in self._sources:
                raise NodeNotFoundError(name)
            state = self._sources[name]
            root = await state.snapshot_service.get_root()
            return replace(_scope_node(name, root), name=name)
        source_name, original_id = _parse_id(parent.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(parent.id)
        state = self._sources[source_name]
        child = await state.snapshot_service.get_child_by_name(name, original_id)
        return _scope_node(source_name, child)

    @override
    async def get_children(self, parent: Node) -> list[Node]:
        if parent.id == _VIRTUAL_ROOT_ID:
            roots: list[Node] = []
            for source_name, state in self._sources.items():
                root = await state.snapshot_service.get_root()
                roots.append(replace(_scope_node(source_name, root), name=source_name))
            return roots
        source_name, original_id = _parse_id(parent.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(parent.id)
        state = self._sources[source_name]
        children = await state.snapshot_service.get_children_by_id(original_id)
        return [_scope_node(source_name, child) for child in children]

    @override
    async def get_trashed_nodes(self, flatten: bool = False) -> list[Node]:
        all_nodes: list[Node] = []
        for source_name, state in self._sources.items():
            trashed = await state.snapshot_service.get_trashed_nodes()
            scoped = [_scope_node(source_name, n) for n in trashed]
            if flatten:
                all_nodes.extend(scoped)
            else:
                ancestor_set = set(n.id for n in scoped if n.is_directory)
                if not ancestor_set:
                    all_nodes.extend(scoped)
                else:
                    table = {n.id: n for n in scoped}
                    all_nodes.extend(
                        n
                        for n in scoped
                        if not _in_ancestor_set(table, n, ancestor_set)
                    )
        return all_nodes

    @override
    async def find_nodes_by_regex(self, pattern: str) -> list[Node]:
        results: list[Node] = []
        for source_name, state in self._sources.items():
            nodes = await state.snapshot_service.find_nodes_by_regex(pattern)
            results.extend(_scope_node(source_name, n) for n in nodes)
        return results

    @override
    async def walk(
        self,
        node: Node,
        *,
        include_trashed: bool = False,
    ) -> AsyncIterator[tuple[Node, list[Node], list[Node]]]:
        if not node.is_directory:
            return
        if node.is_trashed and not include_trashed:
            return

        q = deque([node])
        while q:
            node = q.popleft()
            children = await self.get_children(node)
            children = (_ for _ in children if not _.is_trashed or include_trashed)

            directories, files = tee(children, 2)
            directories = [_ for _ in directories if _.is_directory]
            files = [_ for _ in files if not _.is_directory]

            yield node, directories, files

            q.extend(directories)

    @override
    async def create_directory(
        self,
        name: str,
        parent: Node,
        *,
        exist_ok: bool = False,
    ) -> Node:
        if not parent.is_directory:
            raise ValueError("parent is not a directory")
        if not name:
            raise ValueError("directory name is empty")
        if not is_valid_name(name):
            raise ValueError("no `/` or `\\` allowed in directory name")
        if not await self.is_authenticated():
            raise AuthenticationError()

        if parent.id == _VIRTUAL_ROOT_ID:
            raise ValueError("cannot create directory in virtual root")

        source_name, _ = _parse_id(parent.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(parent.id)
        state = self._sources[source_name]
        unscoped_parent = _unscope_node(source_name, parent)

        if not exist_ok:
            node = await else_none(self.get_child_by_name(name, parent))
            if node:
                raise NodeExistsError(node)

        result = await state.file_service.create_directory(
            name, unscoped_parent, exist_ok=exist_ok, private=None
        )
        return _scope_node(source_name, result)

    @asynccontextmanager
    @override
    async def download_file(self, node: Node) -> AsyncIterator[ReadableFile]:
        if node.is_directory:
            raise ValueError("node should be a file")
        if not await self.is_authenticated():
            raise AuthenticationError()

        source_name, _ = _parse_id(node.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(node.id)
        state = self._sources[source_name]
        unscoped_node = _unscope_node(source_name, node)

        async with state.file_service.download_file(unscoped_node) as file:
            yield file

    @asynccontextmanager
    @override
    async def upload_file(
        self,
        name: str,
        parent: Node,
        *,
        size: int | None = None,
        mime_type: str | None = None,
        media_info: MediaInfo | None = None,
    ) -> AsyncIterator[WritableFile]:
        if not parent.is_directory:
            raise ValueError("parent is not a directory")
        if not name:
            raise ValueError("directory name is empty")
        if not is_valid_name(name):
            raise ValueError("no `/` or `\\` allowed in directory name")
        if not await self.is_authenticated():
            raise AuthenticationError()

        if parent.id == _VIRTUAL_ROOT_ID:
            raise ValueError("cannot upload to virtual root")

        source_name, _ = _parse_id(parent.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(parent.id)
        state = self._sources[source_name]
        unscoped_parent = _unscope_node(source_name, parent)

        node = await else_none(self.get_child_by_name(name, parent))
        if node:
            raise NodeExistsError(node)

        async with state.file_service.upload_file(
            name,
            unscoped_parent,
            size=size,
            mime_type=mime_type,
            media_info=media_info,
            private=None,
        ) as file:
            yield file

    @override
    async def purge_trash(self) -> None:
        if not await self.is_authenticated():
            raise AuthenticationError()
        for state in self._sources.values():
            await state.file_service.purge_trash()

    @override
    async def delete(self, node: Node, *, permanent: bool = False) -> None:
        if not await self.is_authenticated():
            raise AuthenticationError()

        source_name, _ = _parse_id(node.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(node.id)
        state = self._sources[source_name]
        unscoped_node = _unscope_node(source_name, node)
        await state.file_service.delete(unscoped_node, permanent=permanent)

    @override
    async def restore(self, node: Node) -> Node:
        if not await self.is_authenticated():
            raise AuthenticationError()

        source_name, _ = _parse_id(node.id)
        if source_name not in self._sources:
            raise NodeNotFoundError(node.id)
        state = self._sources[source_name]
        unscoped_node = _unscope_node(source_name, node)
        result = await state.file_service.restore(unscoped_node)
        return _scope_node(source_name, result)

    @override
    async def move(
        self,
        node: Node,
        *,
        new_parent: Node | None = None,
        new_name: str | None = None,
    ) -> Node:
        if not await self.is_authenticated():
            raise AuthenticationError()

        if not new_parent and not new_name:
            raise ValueError("nothing to move")

        if new_name and not is_valid_name(new_name):
            raise ValueError("no `/` or `\\` allowed in file name")

        if node.id == _VIRTUAL_ROOT_ID:
            raise ValueError("root node is immutable")

        if node.parent_id == _VIRTUAL_ROOT_ID:
            raise ValueError("source-root node is immutable")

        source_name, _ = _parse_id(node.id)

        if new_parent:
            if new_parent.id == _VIRTUAL_ROOT_ID:
                raise ValueError("cannot move to virtual root")
            new_parent_source, _ = _parse_id(new_parent.id)
            if new_parent_source != source_name:
                raise ValueError("cannot move across sources")
            if not new_parent.is_directory:
                raise ValueError("new_parent is not a directory")
            if await _contains(self, node, new_parent):
                raise ValueError("new_parent is a descendant of the source node")

        if source_name not in self._sources:
            raise NodeNotFoundError(node.id)
        state = self._sources[source_name]
        unscoped_node = _unscope_node(source_name, node)
        unscoped_new_parent = (
            _unscope_node(source_name, new_parent) if new_parent else None
        )

        result = await state.file_service.move(
            unscoped_node,
            new_parent=unscoped_new_parent,
            new_name=new_name,
        )
        return _scope_node(source_name, result)

    @override
    async def sync(self) -> AsyncIterator[ChangeAction]:
        if not await self.is_authenticated():
            raise AuthenticationError()

        queue: Queue[ChangeAction | None] = Queue()
        n_sources = len(self._sources)

        async def sync_source(source_name: str, state: _SourceState) -> None:
            async with state.lock:
                fs = state.file_service
                ss = state.snapshot_service

                initial_cursor = await fs.get_initial_cursor()

                cursor = await ss.get_current_cursor()
                if not cursor:
                    cursor = initial_cursor

                if cursor == initial_cursor:
                    node = await fs.get_root()
                    await ss.set_root(node)

                async for changes, next_ in fs.get_changes(cursor):
                    await ss.apply_changes(changes, next_)

                    for change in changes:
                        await queue.put(_scope_change(source_name, change))
            await queue.put(None)

        async with TaskGroup() as tg:
            for source_name, state in self._sources.items():
                tg.create_task(sync_source(source_name, state))

            done = 0
            while done < n_sources:
                item = await queue.get()
                if item is None:
                    done += 1
                else:
                    yield item

    @override
    async def get_hasher_factory(self) -> CreateHasher:
        first_state = next(iter(self._sources.values()))
        return await first_state.file_service.get_hasher_factory()

    @override
    async def is_authenticated(self) -> bool:
        for state in self._sources.values():
            if not await state.file_service.is_authenticated():
                return False
        return True

    @override
    async def authenticate(self) -> None:
        for state in self._sources.values():
            await state.file_service.authenticate()


def _in_ancestor_set(
    table: dict[str, Node], node: Node, ancestor_set: set[str]
) -> bool:
    if node.parent_id is None:
        return False
    parent = table.get(node.parent_id, None)
    if not parent:
        return False
    if parent.id in ancestor_set:
        return True
    included = _in_ancestor_set(table, parent, ancestor_set)
    if included:
        ancestor_set.add(parent.id)
    return included


async def _contains(drive: Drive, ancestor: Node, node: Node) -> bool:
    visited: set[str] = set()
    while True:
        if ancestor.id == node.id:
            # meet the ancestor
            return True
        if not node.parent_id:
            # reached the root but never meet the ancestor
            return False

        visited.add(node.id)
        node = await drive.get_node_by_id(node.parent_id)
        if node.id in visited:
            raise RuntimeError("detected node cycle")
