"""
NoSQL storage backend for ActivityPub data.
"""

import asyncio
from collections.abc import AsyncGenerator, AsyncIterator, Mapping
import hashlib
from logging import Logger
import math
from typing import Any, cast, override
from uuid import uuid4

from fastapi.datastructures import UploadFile
import gridfs
from pydantic import SecretStr
from pymongo.asynchronous.collection import AsyncCollection
from pymongo.asynchronous.database import AsyncDatabase
from pymongo.asynchronous.mongo_client import AsyncMongoClient

from phederation.models import APDocument, CRUD
from phederation.models.activities import APActivity, APFollow
from phederation.models.actors import APAccount, APActor
from phederation.models.collections import APOrderedCollectionPage
from phederation.models.keys import APPrivateKey
from phederation.models.maintenance import APMaintenanceCommand
from phederation.models.media import MediaMetadata
from phederation.models.objects import APCollectionItem, APObject
from phederation.models.proofs import DataIntegrityProof
from phederation.storage.nosqlcrud import NoSqlCRUD, nosql_build_search_query
from phederation.utils import ObjectId
from phederation.utils.base import (
    AccessType,
    ActivityPubBaseWithId,
    NodeInfo,
    UrlType,
    assemble_id_url,
)
from phederation.utils.exceptions import StorageError, catch_exceptions
from phederation.utils.logging import configure_logger
from phederation.utils.settings import PhedSettings

from .base import StorageBackend, fill_database_url_from_env


class NoSQLStorageBackend(StorageBackend):
    """NoSQL storage backend implementation."""

    def __init__(self, settings: PhedSettings, database_url: SecretStr):
        """Initialize storage backend."""
        super().__init__(database_url=database_url, settings=settings)
        self.logger: Logger = configure_logger(__name__, prefix=self.settings.federation.logging_prefix)

        # Create CRUD objects
        self._crud_tables: dict[str, CRUD[Any]] = {}

        self.actor: CRUD[APActor] = self.table(APActor, name=UrlType.Actors.value)
        self.account: CRUD[APAccount] = self.table(APAccount, name=UrlType.Accounts.value)
        self.activity: CRUD[APActivity] = self.table(APActivity, name=UrlType.Activities.value)
        self.object: CRUD[APObject] = self.table(APObject, name=UrlType.Objects.value)
        self.collection: CRUD[APOrderedCollectionPage] = self.table(APOrderedCollectionPage, name=UrlType.Collections.value)
        self.collection_items: CRUD[APCollectionItem] = self.table(APCollectionItem, name=UrlType.CollectionItems.value)
        self.key: CRUD[APPrivateKey] = self.table(APPrivateKey, name=UrlType.Keys.value)
        self.proof: CRUD[DataIntegrityProof] = self.table(DataIntegrityProof, name=UrlType.Proofs.value)
        self.follow: CRUD[APFollow] = self.table(APFollow, name=UrlType.Follows.value)
        self.media: CRUD[APDocument] = self.table(APDocument, name=UrlType.Media.value)
        self.node_info: CRUD[NodeInfo] = self.table(NodeInfo, name=UrlType.NodeInfo.value)
        self.maintenance: CRUD[APMaintenanceCommand] = self.table(APMaintenanceCommand, name=UrlType.Maintenance.value)

    @override
    def table[T: ActivityPubBaseWithId](self, model_class: type[T], name: str) -> CRUD[T]:
        if not name in self._crud_tables:
            self._crud_tables[name] = NoSqlCRUD[model_class](
                get_collection=self.get_nosql_collection,
                base_url=self.settings.domain.base_url,
                table_name=name,
                logger=configure_logger(f"{__name__}_{name}", prefix=self.settings.federation.logging_prefix),
                MODEL=model_class,
            )
        return self._crud_tables[name]

    async def get_nosql_collection(self, table_name: str) -> AsyncCollection[Any]:
        return (await self.database()).get_collection(table_name)

    @catch_exceptions(StorageError, "Failed to initialize nosql storage")
    @override
    async def initialize(
        self,
    ) -> None:
        """Initialize database connection."""
        self._database: AsyncDatabase[Any] | None = None
        self._client: AsyncMongoClient[Any] | None = None
        self.database_url: SecretStr = fill_database_url_from_env(self.database_url)

    async def database(self):
        loop = asyncio.get_running_loop()
        if self._client is None or self._database is None or self._client._loop != loop or self._client._closed:  # pyright: ignore[reportPrivateUsage]
            try:
                self._client = AsyncMongoClient(self.database_url.get_secret_value())
                self._database = self._client.get_database(self.settings.storage.database.name)
            except Exception as e:
                raise StorageError(f"Could not connect to database; error {type(e).__name__}, {e}")
        return self._database

    @catch_exceptions(StorageError, "NoSql: Failed to reconnect to database")
    @override
    async def reconnect(self) -> None:
        """Reconnect to the database. Useful if the async event loop changed, e.g. in delivery workers in different processes."""
        self._database = None
        _ = await self.database()

    @override
    async def paginate(self, collection_id: ObjectId, access: AccessType = AccessType.PUBLIC) -> AsyncGenerator[list[APCollectionItem], Any]:
        collection_items = await self.get_nosql_collection(await self.collection_items.get_table_name())
        query = nosql_build_search_query(collection_id=collection_id, access=access)
        total_items = await collection_items.count_documents(query)
        n_pages = max(1, math.ceil(total_items / float(self.settings.storage.page_size)))
        for page in range(0, n_pages):
            n_skip = self.settings.storage.page_size * page
            items: list[APCollectionItem] = []
            async with collection_items.find(query).skip(n_skip).limit(self.settings.storage.page_size) as cursor:
                items.extend([APCollectionItem.deserialize(item) async for item in cursor])  # pyright: ignore[reportArgumentType, reportAny]
            yield items

    @override
    async def get_collection_page_items(
        self, collection_id: ObjectId, page: int = 1, access: AccessType = AccessType.PUBLIC
    ) -> list[APCollectionItem] | None:
        collection_items = await self.get_nosql_collection(await self.collection_items.get_table_name())
        n_skip = self.settings.storage.page_size * max(0, page - 1)
        items: list[APCollectionItem] = []

        # construct the search query based on the access level
        query = nosql_build_search_query(collection_id=collection_id, access=access)

        async with collection_items.find(query).skip(n_skip).limit(self.settings.storage.page_size) as cursor:
            async for item in cursor:  # pyright: ignore[reportAny]
                items.append(APCollectionItem.deserialize(item))  # pyright: ignore[reportArgumentType, reportAny]
        return items

    @override
    async def save(self, document_id: ObjectId, content: UploadFile) -> ObjectId:
        database = await self.database()
        files_collection = database.fs.files
        bucket = gridfs.AsyncGridFSBucket(database)
        chunk_size_bytes = 1 * 1024 * 1024
        file_primary = str(uuid4()) + str(hash(document_id))

        file_id = assemble_id_url(
            type=UrlType.Media,
            base_url=self.settings.domain.base_url,
            primary=file_primary,
        )

        file_content_hash = hashlib.sha256()
        async with bucket.open_upload_stream_with_id(
            file_id=file_id,
            filename=file_primary,
            chunk_size_bytes=chunk_size_bytes,
            metadata={"contentType": "application/octet-stream", "filename": file_primary, "documentId": document_id},
        ) as grid_in:
            while True:
                chunk = await content.read(size=chunk_size_bytes)
                if not chunk:
                    break
                await grid_in.write(chunk)
                file_content_hash.update(chunk)

        # update the metadata of the file with the hex digest
        hash_digest = file_content_hash.hexdigest()
        file_data = await files_collection.find_one({"_id": file_id})
        if file_data:
            file_metadata = cast(dict[str, str], file_data.get("metadata", {}))  # pyright: ignore[reportAny]
            file_metadata["sha256"] = hash_digest
            file_metadata["documentId"] = document_id
            _ = await files_collection.update_one({"_id": file_id}, {"$set": {"metadata": file_metadata}})

        return file_id

    @override
    async def load(self, file_id: ObjectId) -> AsyncIterator[bytes]:
        database = await self.database()
        bucket = gridfs.AsyncGridFSBucket(database)

        file = await bucket.open_download_stream(file_id=file_id)
        return file

    @override
    async def metadata(self, file_id: ObjectId) -> MediaMetadata:
        database = await self.database()
        bucket = gridfs.AsyncGridFSBucket(database)

        file = await bucket.open_download_stream(file_id=file_id)
        metadata: Mapping[str, str] | None = file.metadata
        filename = file.name
        content_type: str = "application/json"
        hash_sha256: str | None = None
        if metadata and "contentType" in metadata:
            content_type = metadata.get("contentType", "")
        if metadata and "sha256" in metadata:
            hash_sha256 = metadata.get("sha256", "")
        if metadata and "documentId" in metadata:
            document_id = metadata.get("documentId", "")
        else:
            raise StorageError(f"documentId not in metadata of file {file_id}")
        return MediaMetadata(content_type=content_type, filename=filename, hash_sha256=hash_sha256, document_id=document_id)

    @override
    async def is_following(self, follower: str, following: str, check_accepted: bool = False) -> bool:
        """Check if follower 'follower' is following 'following'."""
        try:
            database = await self.database()
            collection = database.get_collection(await self.follow.get_table_name())
            if check_accepted:
                follow = await collection.find_one({"actor": follower, "object": following, "accepted": True})
            else:
                follow = await collection.find_one({"actor": follower, "object": following})
            return follow is not None

        except Exception as e:
            raise StorageError(f"NoSQLStorageBackend: Failed to check if follower '{follower}' is following '{following}': {e}")

    @override
    async def flush(self) -> None:
        """Flush storage to file, if possible."""
        pass

    @override
    async def clear(self):
        """Delete the database. Useful for testing."""
        # drop all collections at the start, to clean up
        database = await self.database()
        collection_names = await database.list_collection_names()
        for collection_name in collection_names:
            collection = database[collection_name]
            await collection.drop()

    @override
    async def close(self) -> None:
        """Close database connection."""
        pass
