import typing as t
from contextlib import asynccontextmanager
from functools import cached_property

from google.cloud import firestore
from google.cloud.firestore_v1.base_query import FieldFilter
from acb.config import Config
from acb.depends import depends

from ._base import NosqlBase, NosqlBaseSettings


class NosqlSettings(NosqlBaseSettings):
    project_id: t.Optional[str] = None
    credentials_path: t.Optional[str] = None
    emulator_host: t.Optional[str] = None

    @depends.inject
    def __init__(self, config: Config = depends(), **values: t.Any) -> None:
        super().__init__(**values)
        if not self.project_id:
            self.project_id = config.app.project


class Nosql(NosqlBase):
    _transaction = None

    @cached_property
    def client(self) -> firestore.Client:
        if not self._client:
            kwargs = {}
            if self.config.nosql.project_id:
                kwargs["project"] = self.config.nosql.project_id
            if self.config.nosql.credentials_path:
                kwargs["credentials"] = self.config.nosql.credentials_path

            self._client = firestore.Client(**kwargs)
        return self._client

    @cached_property
    def db(self) -> firestore.Client:
        return self.client

    async def init(self) -> None:
        self.logger.info(
            f"Initializing Firestore connection for project {self.config.nosql.project_id}"
        )
        try:
            self.client.collection("test")
            self.logger.info("Firestore connection initialized successfully")
        except Exception as e:
            self.logger.error(f"Failed to initialize Firestore connection: {e}")
            raise

    def _get_collection_ref(self, collection: str) -> firestore.CollectionReference:
        prefix = self.config.nosql.collection_prefix
        return self.client.collection(f"{prefix}{collection}")

    def _convert_to_dict(self, doc: firestore.DocumentSnapshot) -> t.Dict[str, t.Any]:
        if not doc.exists:
            return {}

        data = doc.to_dict()
        if data is None:
            data = {}
        data["_id"] = doc.id
        return data

    def _prepare_document(self, document: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
        if "_id" in document:
            document = document.copy()
            del document["_id"]
        return document

    async def find(
        self, collection: str, filter: t.Dict[str, t.Any], **kwargs: t.Any
    ) -> t.List[t.Dict[str, t.Any]]:
        collection_ref = self._get_collection_ref(collection)
        query = collection_ref

        for key, value in filter.items():
            if key == "_id":
                continue
            query = query.where(filter=FieldFilter(key, "==", value))

        if "limit" in kwargs:
            query = query.limit(kwargs["limit"])

        if "order_by" in kwargs:
            for field in kwargs["order_by"]:
                direction = (
                    firestore.Query.DESCENDING
                    if field.startswith("-")
                    else firestore.Query.ASCENDING
                )
                field_name = field.removeprefix("-")
                query = query.order_by(field_name, direction=direction)

        docs = query.stream()

        results = [self._convert_to_dict(doc) for doc in docs]

        if "_id" in filter:
            results = [doc for doc in results if doc.get("_id") == filter["_id"]]

        return results

    async def find_one(
        self, collection: str, filter: t.Dict[str, t.Any], **kwargs: t.Any
    ) -> t.Optional[t.Dict[str, t.Any]]:
        if "_id" in filter:
            doc_ref = self._get_collection_ref(collection).document(filter["_id"])
            doc = doc_ref.get()
            if doc.exists:
                return self._convert_to_dict(doc)
            return None

        results = await self.find(collection, filter, limit=1, **kwargs)
        return results[0] if results else None

    async def insert_one(
        self, collection: str, document: t.Dict[str, t.Any], **kwargs: t.Any
    ) -> t.Any:
        collection_ref = self._get_collection_ref(collection)

        if "_id" in document:
            doc_id = document["_id"]
            doc_ref = collection_ref.document(doc_id)
            doc_ref.set(self._prepare_document(document))
        else:
            doc_ref = collection_ref.add(document)[1]
            doc_id = doc_ref.id

        return doc_id

    async def insert_many(
        self, collection: str, documents: t.List[t.Dict[str, t.Any]], **kwargs: t.Any
    ) -> t.List[t.Any]:
        ids = []
        batch = self.client.batch()
        collection_ref = self._get_collection_ref(collection)

        for doc in documents:
            if "_id" in doc:
                doc_id = doc["_id"]
                doc_ref = collection_ref.document(doc_id)
                batch.set(doc_ref, self._prepare_document(doc))
            else:
                doc_ref = collection_ref.document()
                batch.set(doc_ref, doc)
                doc_id = doc_ref.id

            ids.append(doc_id)

        batch.commit()
        return ids

    async def update_one(
        self,
        collection: str,
        filter: t.Dict[str, t.Any],
        update: t.Dict[str, t.Any],
        **kwargs: t.Any,
    ) -> t.Any:
        doc = await self.find_one(collection, filter)
        if not doc:
            return {"modified_count": 0}

        doc_id = doc["_id"]
        doc_ref = self._get_collection_ref(collection).document(doc_id)

        if "$set" in update:
            doc_ref.update(update["$set"])
        else:
            doc_ref.update(update)

        return {"modified_count": 1}

    async def update_many(
        self,
        collection: str,
        filter: t.Dict[str, t.Any],
        update: t.Dict[str, t.Any],
        **kwargs: t.Any,
    ) -> t.Any:
        docs = await self.find(collection, filter)
        if not docs:
            return {"modified_count": 0}

        batch = self.client.batch()
        collection_ref = self._get_collection_ref(collection)

        for doc in docs:
            doc_id = doc["_id"]
            doc_ref = collection_ref.document(doc_id)

            if "$set" in update:
                batch.update(doc_ref, update["$set"])
            else:
                batch.update(doc_ref, update)

        batch.commit()
        return {"modified_count": len(docs)}

    async def delete_one(
        self, collection: str, filter: t.Dict[str, t.Any], **kwargs: t.Any
    ) -> t.Any:
        doc = await self.find_one(collection, filter)
        if not doc:
            return {"deleted_count": 0}

        doc_id = doc["_id"]
        doc_ref = self._get_collection_ref(collection).document(doc_id)

        doc_ref.delete()

        return {"deleted_count": 1}

    async def delete_many(
        self, collection: str, filter: t.Dict[str, t.Any], **kwargs: t.Any
    ) -> t.Any:
        docs = await self.find(collection, filter)
        if not docs:
            return {"deleted_count": 0}

        batch = self.client.batch()
        collection_ref = self._get_collection_ref(collection)

        for doc in docs:
            doc_id = doc["_id"]
            doc_ref = collection_ref.document(doc_id)
            batch.delete(doc_ref)

        batch.commit()
        return {"deleted_count": len(docs)}

    async def count(
        self,
        collection: str,
        filter: t.Optional[t.Dict[str, t.Any]] = None,
        **kwargs: t.Any,
    ) -> int:
        docs = await self.find(collection, filter or {})
        return len(docs)

    async def aggregate(
        self, collection: str, pipeline: t.List[t.Dict[str, t.Any]], **kwargs: t.Any
    ) -> t.List[t.Dict[str, t.Any]]:
        docs = await self.find(collection, {})

        for stage in pipeline:
            if "$match" in stage:
                filter_dict = stage["$match"]
                docs = [
                    doc
                    for doc in docs
                    if all(doc.get(k) == v for k, v in filter_dict.items())
                ]
            elif "$project" in stage:
                projection = stage["$project"]
                docs = [{k: doc[k] for k in projection if k in doc} for doc in docs]
            elif "$limit" in stage:
                docs = docs[: stage["$limit"]]
            elif "$skip" in stage:
                docs = docs[stage["$skip"] :]

        return docs

    @asynccontextmanager
    async def transaction(self) -> t.AsyncGenerator[None, None]:
        transaction = self.client.transaction()
        try:
            self._transaction = transaction
            with transaction:
                yield
        except Exception as e:
            self.logger.error(f"Transaction failed: {e}")
            raise
        finally:
            self._transaction = None


depends.set(Nosql)
