from datetime import datetime, timezone
from typing import Dict, Generic, List, Optional, Type, TypeVar

from bson import ObjectId
from motor.motor_asyncio import (
    AsyncIOMotorClient,
    AsyncIOMotorClientSession,
    AsyncIOMotorDatabase,
)
from pydantic import BaseModel

from .audit import log_audit
from .exceptions import DocumentNotFoundException

T = TypeVar("T", bound=BaseModel)


class LazyQL(Generic[T]):
    """
    Generic CRUD service with built-in audit logging and ACID transactions.
    """

    def __init__(
        self,
        db: AsyncIOMotorDatabase,
        collection_name: str,
        model_cls: Type[T],
        client: Optional[AsyncIOMotorClient] = None,
    ):
        self.db = db
        self.collection = db[collection_name]
        self.model_cls = model_cls
        self.collection_name = collection_name
        self.client = client

    async def get_all(
        self,
        limit: Optional[int] = 100,
        skip: Optional[int] = 0,
        filters: Optional[Dict] = None,
        sort_by: Optional[str] = None,
        sort_order: int = 1,
        include_deleted: bool = False,
    ) -> List[T]:
        """
        Get all documents with optional filtering, pagination, sorting and
        soft-delete handling.
        """
        query = filters or {}

        # Soft Delete: Exclude documents with deleted_at set
        if not include_deleted:
            query["deleted_at"] = None

        cursor = self.collection.find(query)

        if skip is not None:
            cursor = cursor.skip(skip)
        if limit is not None:
            cursor = cursor.limit(limit)

        if sort_by:
            cursor = cursor.sort(sort_by, sort_order)

        docs = [self.model_cls(**doc) async for doc in cursor]
        return docs

    async def get_one(self, doc_id: str) -> Optional[T]:
        doc = await self.collection.find_one(
            {"_id": ObjectId(doc_id), "deleted_at": None}
        )
        if not doc:
            return None
        return self.model_cls(**doc)

    async def create(self, data: BaseModel, user: str = "system") -> T:
        if self.client:
            async with await self.client.start_session() as session:
                async with session.start_transaction():
                    return await self._create_implementation(
                        data, user, session
                    )
        else:
            return await self._create_implementation(data, user, None)

    async def _create_implementation(
        self,
        data: BaseModel,
        user: str,
        session: Optional[AsyncIOMotorClientSession],
    ) -> T:
        doc = data.model_dump(by_alias=True, exclude={"id"})

        now = datetime.now(timezone.utc)
        doc.update(
            {
                "created_at": now,
                "updated_at": now,
                "created_by": user,
                "updated_by": user,
                "deleted_at": None,
            }
        )

        result = await self.collection.insert_one(doc, session=session)
        doc["_id"] = result.inserted_id

        await log_audit(
            self.db,
            self.collection_name,
            result.inserted_id,
            user,
            "CREATE",
            session=session,
        )

        return self.model_cls(**doc)

    async def update(self, id: str, data: BaseModel, user: str) -> T:
        if self.client:
            async with await self.client.start_session() as session:
                async with session.start_transaction():
                    return await self._update_implementation(
                        id, data, user, session
                    )
        else:
            return await self._update_implementation(id, data, user, None)

    async def _update_implementation(
        self,
        id: str,
        data: BaseModel,
        user: str,
        session: Optional[AsyncIOMotorClientSession],
    ) -> T:
        # Fetch old document for audit logging
        old_doc = await self.collection.find_one(
            {"_id": ObjectId(id), "deleted_at": None}, session=session
        )
        if not old_doc:
            raise DocumentNotFoundException()

        update_data = data.model_dump(exclude_unset=True, exclude={"id"})
        update_data.update(
            {"updated_at": datetime.now(timezone.utc), "updated_by": user}
        )

        new_doc = await self.collection.find_one_and_update(
            {"_id": ObjectId(id)},
            {"$set": update_data},
            return_document=True,
            session=session,
        )

        await log_audit(
            self.db,
            self.collection_name,
            id,
            user,
            "UPDATE",
            old_doc,
            new_doc,
            session=session,
        )

        return self.model_cls(**new_doc)

    async def delete(self, id: str, user: str) -> None:
        """
        Soft delete a document (sets deleted_at timestamp).
        """
        if self.client:
            async with await self.client.start_session() as session:
                async with session.start_transaction():
                    await self._delete_implementation(id, user, session)
        else:
            await self._delete_implementation(id, user, None)

    async def _delete_implementation(
        self, id: str, user: str, session: Optional[AsyncIOMotorClientSession]
    ) -> None:
        old_doc = await self.collection.find_one(
            {"_id": ObjectId(id), "deleted_at": None}, session=session
        )
        if not old_doc:
            raise DocumentNotFoundException()

        await self.collection.update_one(
            {"_id": ObjectId(id)},
            {
                "$set": {
                    "deleted_at": datetime.now(timezone.utc),
                    "updated_by": user,
                }
            },
            session=session,
        )

        await log_audit(
            self.db,
            self.collection_name,
            id,
            user,
            "DELETE",
            old_doc,
            None,
            session=session,
        )

    async def hard_delete(self, id: str, user: str) -> None:
        """
        Permanently delete a document from the collection.
        """
        if self.client:
            async with await self.client.start_session() as session:
                async with session.start_transaction():
                    await self._hard_delete_implementation(id, user, session)
        else:
            await self._hard_delete_implementation(id, user, None)

    async def _hard_delete_implementation(
        self, id: str, user: str, session: Optional[AsyncIOMotorClientSession]
    ) -> None:
        await self.collection.delete_one(
            {"_id": ObjectId(id)}, session=session
        )

        await log_audit(
            self.db,
            self.collection_name,
            id,
            user,
            "HARD_DELETE",
            session=session,
        )

    async def restore(self, id: str, user: str) -> T:
        """
        Restore a soft-deleted document (sets deleted_at to None).
        """
        if self.client:
            async with await self.client.start_session() as session:
                async with session.start_transaction():
                    res = await self._restore_implementation(id, user, session)
        else:
            res = await self._restore_implementation(id, user, None)
        return res

    async def _restore_implementation(
        self, id: str, user: str, session: Optional[AsyncIOMotorClientSession]
    ) -> T:
        old_doc = await self.collection.find_one(
            {"_id": ObjectId(id), "deleted_at": {"$ne": None}}, session=session
        )
        if not old_doc:
            raise DocumentNotFoundException()

        new_doc = await self.collection.find_one_and_update(
            {"_id": ObjectId(id)},
            {"$set": {"deleted_at": None, "updated_by": user}},
            return_document=True,
            session=session,
        )

        await log_audit(
            self.db,
            self.collection_name,
            id,
            user,
            "RESTORE",
            old_doc,
            new_doc,
            session=session,
        )
        return self.model_cls(**new_doc)
