from __future__ import annotations

import copy
import re as _re
import uuid
from typing import TYPE_CHECKING, Any, Optional

import pymongo

from mockmesh.core.fallback import FallbackConfig
from mockmesh.core.logger import get_logger
from mockmesh.plugins.base import BasePlugin

if TYPE_CHECKING:
    from mockmesh.response.engine import ResponseEngine
    from mockmesh.storage.manager import StorageManager

_log = get_logger("interceptors.mongodb")


def _match_filter(doc: dict, filter: Optional[dict]) -> bool:
    if not filter:
        return True
    for key, value in filter.items():
        if key == "$and":
            if not all(_match_filter(doc, sub) for sub in value):
                return False
            continue
        if key == "$or":
            if not any(_match_filter(doc, sub) for sub in value):
                return False
            continue
        doc_val = doc.get(key)
        if isinstance(value, dict):
            for op, operand in value.items():
                if op == "$eq":
                    if doc_val != operand:
                        return False
                elif op == "$ne":
                    if doc_val == operand:
                        return False
                elif op == "$gt":
                    if doc_val is None or doc_val <= operand:
                        return False
                elif op == "$gte":
                    if doc_val is None or doc_val < operand:
                        return False
                elif op == "$lt":
                    if doc_val is None or doc_val >= operand:
                        return False
                elif op == "$lte":
                    if doc_val is None or doc_val > operand:
                        return False
                elif op == "$in":
                    if doc_val not in operand:
                        return False
                elif op == "$nin":
                    if doc_val in operand:
                        return False
                elif op == "$exists":
                    if operand and key not in doc:
                        return False
                    if not operand and key in doc:
                        return False
                elif op == "$regex":
                    if doc_val is None or not _re.search(operand, str(doc_val)):
                        return False
        else:
            if doc_val != value:
                return False
    return True


def _apply_update(doc: dict, update: dict) -> bool:
    modified = False
    if "$set" in update:
        for k, v in update["$set"].items():
            if doc.get(k) != v:
                doc[k] = v
                modified = True
    if "$unset" in update:
        for k in update["$unset"]:
            if k in doc:
                del doc[k]
                modified = True
    if "$inc" in update:
        for k, v in update["$inc"].items():
            doc[k] = doc.get(k, 0) + v
            modified = True
    return modified


class _MockInsertOneResult:
    def __init__(self, inserted_id: Any) -> None:
        self.inserted_id = inserted_id
        self.acknowledged = True


class _MockInsertManyResult:
    def __init__(self, inserted_ids: list) -> None:
        self.inserted_ids = inserted_ids
        self.acknowledged = True


class _MockUpdateResult:
    def __init__(self, matched_count: int, modified_count: int) -> None:
        self.matched_count = matched_count
        self.modified_count = modified_count
        self.acknowledged = True


class _MockDeleteResult:
    def __init__(self, deleted_count: int) -> None:
        self.deleted_count = deleted_count
        self.acknowledged = True


class _MockCursor:
    def __init__(self, documents: list) -> None:
        self._documents = list(documents)
        self._skip_count = 0
        self._limit_count = 0

    def __iter__(self):
        docs = self._documents
        if self._skip_count:
            docs = docs[self._skip_count:]
        if self._limit_count:
            docs = docs[:self._limit_count]
        yield from docs

    def sort(self, key_or_list: Any, direction: Optional[int] = None) -> "_MockCursor":
        if isinstance(key_or_list, str):
            reverse = direction == -1 if direction is not None else False
            self._documents.sort(key=lambda d: d.get(key_or_list, ""), reverse=reverse)
        elif isinstance(key_or_list, list):
            for key, d in reversed(key_or_list):
                reverse = d == -1
                self._documents.sort(key=lambda doc, k=key: doc.get(k, ""), reverse=reverse)
        return self

    def limit(self, count: int) -> "_MockCursor":
        self._limit_count = count
        return self

    def skip(self, count: int) -> "_MockCursor":
        self._skip_count = count
        return self

    def count(self) -> int:
        return len(list(self.__iter__()))

    def __len__(self) -> int:
        return self.count()


class _MockCollection:
    def __init__(self, db_name: str, coll_name: str, sm: "StorageManager") -> None:
        self._db_name = db_name
        self._coll_name = coll_name
        self._sm = sm
        self._collection_key = f"mongo--{db_name}--{coll_name}"

    @property
    def name(self) -> str:
        return self._coll_name

    def _all_docs(self) -> list[dict]:
        records = self._sm.nosql.list(self._collection_key)
        return [r["value"] for r in records]

    def _filtered(self, filter: Optional[dict]) -> list[dict]:
        return [d for d in self._all_docs() if _match_filter(d, filter)]

    def insert_one(self, document: dict) -> _MockInsertOneResult:
        doc = dict(document)
        if "_id" not in doc:
            doc["_id"] = str(uuid.uuid4())
        self._sm.nosql.put(self._collection_key, str(doc["_id"]), doc)
        _log.info("MongoDB insert_one", extra={"mm_collection": self._collection_key})
        return _MockInsertOneResult(doc["_id"])

    def insert_many(self, documents: list) -> _MockInsertManyResult:
        ids = []
        for document in documents:
            doc = dict(document)
            if "_id" not in doc:
                doc["_id"] = str(uuid.uuid4())
            self._sm.nosql.put(self._collection_key, str(doc["_id"]), doc)
            ids.append(doc["_id"])
        _log.info("MongoDB insert_many", extra={"mm_collection": self._collection_key, "mm_count": len(ids)})
        return _MockInsertManyResult(ids)

    def find_one(self, filter: Optional[dict] = None) -> Optional[dict]:
        matches = self._filtered(filter)
        return copy.deepcopy(matches[0]) if matches else None

    def find(self, filter: Optional[dict] = None) -> _MockCursor:
        return _MockCursor([copy.deepcopy(d) for d in self._filtered(filter)])

    def update_one(self, filter: dict, update: dict) -> _MockUpdateResult:
        matches = self._filtered(filter)
        if not matches:
            return _MockUpdateResult(0, 0)
        doc = matches[0]
        modified = _apply_update(doc, update)
        self._sm.nosql.put(self._collection_key, str(doc["_id"]), doc)
        return _MockUpdateResult(1, 1 if modified else 0)

    def update_many(self, filter: dict, update: dict) -> _MockUpdateResult:
        matches = self._filtered(filter)
        modified_count = 0
        for doc in matches:
            if _apply_update(doc, update):
                modified_count += 1
            self._sm.nosql.put(self._collection_key, str(doc["_id"]), doc)
        return _MockUpdateResult(len(matches), modified_count)

    def delete_one(self, filter: dict) -> _MockDeleteResult:
        matches = self._filtered(filter)
        if not matches:
            return _MockDeleteResult(0)
        self._sm.nosql.delete(self._collection_key, str(matches[0]["_id"]))
        return _MockDeleteResult(1)

    def delete_many(self, filter: dict) -> _MockDeleteResult:
        matches = self._filtered(filter)
        for doc in matches:
            self._sm.nosql.delete(self._collection_key, str(doc["_id"]))
        return _MockDeleteResult(len(matches))

    def count_documents(self, filter: dict) -> int:
        return len(self._filtered(filter))

    def replace_one(self, filter: dict, replacement: dict) -> _MockUpdateResult:
        matches = self._filtered(filter)
        if not matches:
            return _MockUpdateResult(0, 0)
        doc = matches[0]
        new_doc = dict(replacement)
        new_doc["_id"] = doc["_id"]
        self._sm.nosql.put(self._collection_key, str(doc["_id"]), new_doc)
        return _MockUpdateResult(1, 1)

    def create_index(self, keys: Any, **kwargs: Any) -> str:
        if isinstance(keys, str):
            return f"{keys}_1"
        if isinstance(keys, list):
            parts = []
            for k, d in keys:
                parts.append(f"{k}_{d}")
            return "_".join(parts)
        return "index_0"

    def aggregate(self, pipeline: list) -> _MockCursor:
        docs = self._all_docs()
        for stage in pipeline:
            if "$match" in stage:
                docs = [d for d in docs if _match_filter(d, stage["$match"])]
            elif "$project" in stage:
                projection = stage["$project"]
                projected = []
                for d in docs:
                    new_doc: dict[str, Any] = {}
                    for field, include in projection.items():
                        if include and field in d:
                            new_doc[field] = d[field]
                    if "_id" in d and projection.get("_id", 1):
                        new_doc["_id"] = d["_id"]
                    projected.append(new_doc)
                docs = projected
        return _MockCursor([copy.deepcopy(d) for d in docs])


class _MockDatabase:
    def __init__(self, db_name: str, sm: "StorageManager") -> None:
        self._db_name = db_name
        self._sm = sm
        self._collections: dict[str, _MockCollection] = {}

    @property
    def name(self) -> str:
        return self._db_name

    def _get_collection(self, name: str) -> _MockCollection:
        if name not in self._collections:
            self._collections[name] = _MockCollection(self._db_name, name, self._sm)
        return self._collections[name]

    def __getitem__(self, name: str) -> _MockCollection:
        return self._get_collection(name)

    def __getattr__(self, name: str) -> _MockCollection:
        if name.startswith("_"):
            raise AttributeError(name)
        return self._get_collection(name)

    def list_collection_names(self) -> list[str]:
        prefix = f"mongo--{self._db_name}--"
        all_colls = self._sm.nosql.list_collections()
        return [c[len(prefix):] for c in all_colls if c.startswith(prefix)]

    def create_collection(self, name: str) -> _MockCollection:
        return self._get_collection(name)

    def drop_collection(self, name: str) -> None:
        coll_key = f"mongo--{self._db_name}--{name}"
        for record in self._sm.nosql.list(coll_key):
            self._sm.nosql.delete(coll_key, record["key"])
        self._collections.pop(name, None)

    def command(self, cmd: Any, **kwargs: Any) -> dict:
        return {"ok": 1}


class _MockMongoClient:
    def __init__(self, sm: "StorageManager", *args: Any, **kwargs: Any) -> None:
        self._sm = sm
        self._databases: dict[str, _MockDatabase] = {}

    def _get_database(self, name: str) -> _MockDatabase:
        if name not in self._databases:
            self._databases[name] = _MockDatabase(name, self._sm)
        return self._databases[name]

    def __getitem__(self, name: str) -> _MockDatabase:
        return self._get_database(name)

    def __getattr__(self, name: str) -> _MockDatabase:
        if name.startswith("_"):
            raise AttributeError(name)
        return self._get_database(name)

    def list_database_names(self) -> list[str]:
        all_colls = self._sm.nosql.list_collections()
        dbs: set[str] = set()
        for c in all_colls:
            if c.startswith("mongo--"):
                parts = c.split("--", 2)
                if len(parts) >= 2:
                    dbs.add(parts[1])
        return sorted(dbs)

    def close(self) -> None:
        pass

    def __enter__(self) -> "_MockMongoClient":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()


class MongoDBInterceptor(BasePlugin):

    name = "mongodb"

    def __init__(self, response_engine: "ResponseEngine") -> None:
        super().__init__(response_engine)
        self._original_client: Any = None

    def activate(
        self,
        config: dict,
        *,
        storage: "Any | None" = None,
        provider_defaults: Optional[dict] = None,
        fallback_config: Optional[FallbackConfig] = None,
        custom_handlers: Optional[dict] = None,
        **_: Any,
    ) -> None:
        self._fallback_config = fallback_config
        self._custom_handlers = custom_handlers if custom_handlers is not None else {}
        self.storage = storage
        self._defaults: dict = {}
        if provider_defaults is not None:
            self._defaults = provider_defaults.get("mongodb", {})
        self._patch_pymongo()
        self._active = True
        _log.info("MongoDB interceptor activated")

    def deactivate(self) -> None:
        if not self._active:
            return
        self._unpatch_pymongo()
        self._active = False
        _log.info("MongoDB interceptor deactivated")

    def _patch_pymongo(self) -> None:
        sm = self.storage
        self._original_client = pymongo.MongoClient

        class _Client(_MockMongoClient):
            def __init__(self, *args: Any, **kwargs: Any) -> None:
                super().__init__(sm, *args, **kwargs)

        pymongo.MongoClient = _Client  # type: ignore
        _log.debug("Patched pymongo.MongoClient")

    def _unpatch_pymongo(self) -> None:
        if self._original_client is None:
            return
        pymongo.MongoClient = self._original_client
        self._original_client = None
        _log.debug("Restored pymongo.MongoClient")
