"""pymongo replay for AutoDebug.

This module provides replay of captured MongoDB operations during test replay.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from autodebug_replay.transcripts import TranscriptIndex

# Original functions
_original_find: Optional[Callable] = None
_original_find_one: Optional[Callable] = None
_original_count_documents: Optional[Callable] = None
_original_aggregate: Optional[Callable] = None
_original_distinct: Optional[Callable] = None
_original_write_ops: Dict[str, Callable] = {}

# Global transcript index
_index: Optional["TranscriptIndex"] = None


def _deserialize_value(value: Any) -> Any:
    """Deserialize a recorded value back to Python type."""
    if value is None:
        return None
    if isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, dict):
        # Handle special BSON types
        if "$oid" in value:
            try:
                from bson import ObjectId
                return ObjectId(value["$oid"])
            except ImportError:
                return value["$oid"]
        if "$date" in value:
            try:
                return datetime.fromisoformat(value["$date"].replace("Z", "+00:00"))
            except ValueError:
                return value["$date"]
        if "$binary" in value:
            return bytes.fromhex(value["$binary"])
        # Regular dict
        return {k: _deserialize_value(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_deserialize_value(item) for item in value]
    return value


def _deserialize_document(doc: Any) -> Any:
    """Deserialize a recorded document."""
    return _deserialize_value(doc)


class _ReplayCursor:
    """Mock cursor that returns recorded results."""

    def __init__(self, results: List[Dict]):
        self._results = [_deserialize_document(r) for r in results]
        self._index = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self._index >= len(self._results):
            raise StopIteration
        result = self._results[self._index]
        self._index += 1
        return result

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def to_list(self, length=None):
        """Return results as a list."""
        if length is None:
            return self._results
        return self._results[:length]

    def limit(self, limit):
        """Return a limited cursor."""
        return _ReplayCursor(self._results[:limit])

    def skip(self, skip):
        """Return a cursor with skipped documents."""
        return _ReplayCursor(self._results[skip:])

    def sort(self, key_or_list, direction=None):
        """Return self (sort is already applied in recorded results)."""
        return self

    def count(self):
        """Return count of results."""
        return len(self._results)


def _patched_find(self, filter=None, *args, **kwargs):
    """Patched Collection.find that replays from transcript."""
    from autodebug_replay.transcripts import ReplayDivergence

    record = _index.expect("MONGODB")

    # Verify operation matches
    if record["operation"] != "find":
        raise ReplayDivergence("MONGODB_OPERATION_MISMATCH", {
            "expected": record["operation"],
            "actual": "find",
        })

    # Verify collection matches
    if record["collection"] != self.name:
        raise ReplayDivergence("MONGODB_COLLECTION_MISMATCH", {
            "expected": record["collection"],
            "actual": self.name,
        })

    return _ReplayCursor(record.get("results", []))


def _patched_find_one(self, filter=None, *args, **kwargs):
    """Patched Collection.find_one that replays from transcript."""
    from autodebug_replay.transcripts import ReplayDivergence

    record = _index.expect("MONGODB")

    if record["operation"] != "find_one":
        raise ReplayDivergence("MONGODB_OPERATION_MISMATCH", {
            "expected": record["operation"],
            "actual": "find_one",
        })

    if record["collection"] != self.name:
        raise ReplayDivergence("MONGODB_COLLECTION_MISMATCH", {
            "expected": record["collection"],
            "actual": self.name,
        })

    result = record.get("result")
    return _deserialize_document(result) if result else None


def _patched_count_documents(self, filter, *args, **kwargs):
    """Patched Collection.count_documents that replays from transcript."""
    from autodebug_replay.transcripts import ReplayDivergence

    record = _index.expect("MONGODB")

    if record["operation"] != "count_documents":
        raise ReplayDivergence("MONGODB_OPERATION_MISMATCH", {
            "expected": record["operation"],
            "actual": "count_documents",
        })

    if record["collection"] != self.name:
        raise ReplayDivergence("MONGODB_COLLECTION_MISMATCH", {
            "expected": record["collection"],
            "actual": self.name,
        })

    return record.get("count", 0)


def _patched_aggregate(self, pipeline, *args, **kwargs):
    """Patched Collection.aggregate that replays from transcript."""
    from autodebug_replay.transcripts import ReplayDivergence

    record = _index.expect("MONGODB")

    if record["operation"] != "aggregate":
        raise ReplayDivergence("MONGODB_OPERATION_MISMATCH", {
            "expected": record["operation"],
            "actual": "aggregate",
        })

    if record["collection"] != self.name:
        raise ReplayDivergence("MONGODB_COLLECTION_MISMATCH", {
            "expected": record["collection"],
            "actual": self.name,
        })

    return _ReplayCursor(record.get("results", []))


def _patched_distinct(self, key, filter=None, *args, **kwargs):
    """Patched Collection.distinct that replays from transcript."""
    from autodebug_replay.transcripts import ReplayDivergence

    record = _index.expect("MONGODB")

    if record["operation"] != "distinct":
        raise ReplayDivergence("MONGODB_OPERATION_MISMATCH", {
            "expected": record["operation"],
            "actual": "distinct",
        })

    if record["collection"] != self.name:
        raise ReplayDivergence("MONGODB_COLLECTION_MISMATCH", {
            "expected": record["collection"],
            "actual": self.name,
        })

    return [_deserialize_value(v) for v in record.get("results", [])]


def _create_write_blocker(operation: str) -> Callable:
    """Create a blocker for write operations during replay."""
    def blocker(self, *args, **kwargs):
        from autodebug_replay.transcripts import ReplayDivergence
        raise ReplayDivergence("MONGODB_WRITE_NOT_ALLOWED", {
            "operation": operation,
            "collection": self.name,
            "message": f"MongoDB {operation} cannot be replayed (write operations blocked)",
        })
    return blocker


def install(index: "TranscriptIndex") -> None:
    """Install pymongo replay patches.

    Args:
        index: TranscriptIndex containing recorded MongoDB calls
    """
    global _original_find, _original_find_one, _original_count_documents
    global _original_aggregate, _original_distinct, _index

    try:
        from pymongo.collection import Collection
    except ImportError:
        return

    if _original_find is not None:
        return  # Already patched

    _index = index

    # Patch read operations
    _original_find = Collection.find
    Collection.find = _patched_find

    _original_find_one = Collection.find_one
    Collection.find_one = _patched_find_one

    _original_count_documents = Collection.count_documents
    Collection.count_documents = _patched_count_documents

    _original_aggregate = Collection.aggregate
    Collection.aggregate = _patched_aggregate

    _original_distinct = Collection.distinct
    Collection.distinct = _patched_distinct

    # Block write operations (but preserve originals for clean uninstall)
    write_ops = [
        "insert_one", "insert_many",
        "update_one", "update_many",
        "delete_one", "delete_many",
        "replace_one", "bulk_write",
        "find_one_and_update", "find_one_and_replace", "find_one_and_delete",
    ]
    for op in write_ops:
        if hasattr(Collection, op):
            original = getattr(Collection, op, None)
            if callable(original):
                _original_write_ops[op] = original
                setattr(Collection, op, _create_write_blocker(op))


def uninstall() -> None:
    """Restore original pymongo functions."""
    global _original_find, _original_find_one, _original_count_documents
    global _original_aggregate, _original_distinct, _index

    try:
        from pymongo.collection import Collection
    except ImportError:
        return

    if _original_find is not None:
        Collection.find = _original_find
        _original_find = None

    if _original_find_one is not None:
        Collection.find_one = _original_find_one
        _original_find_one = None

    if _original_count_documents is not None:
        Collection.count_documents = _original_count_documents
        _original_count_documents = None

    if _original_aggregate is not None:
        Collection.aggregate = _original_aggregate
        _original_aggregate = None

    if _original_distinct is not None:
        Collection.distinct = _original_distinct
        _original_distinct = None

    # Restore write operations
    for op, original in _original_write_ops.items():
        if hasattr(Collection, op):
            setattr(Collection, op, original)
    _original_write_ops.clear()

    _index = None
