"""pymongo patching for AutoDebug capture.

This module captures MongoDB operations at the Collection level,
initially supporting read operations (find, find_one, count_documents, aggregate).
"""

from __future__ import annotations

import hashlib
import json
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from autodebug_probe.recorder import Recorder

# 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

# Global recorder reference
_recorder: Optional["Recorder"] = None

# Supported read operations
SUPPORTED_READ_OPS = {
    "find",
    "find_one",
    "count_documents",
    "aggregate",
    "distinct",
}

# Unsupported write operations
UNSUPPORTED_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",
}


def _get_recorder() -> "Recorder":
    """Get the global recorder instance."""
    if _recorder is None:
        from autodebug_probe.recorder import get_recorder
        return get_recorder()
    return _recorder


def _serialize_document(doc: Any) -> Any:
    """Serialize a MongoDB document for recording."""
    if doc is None:
        return None
    if isinstance(doc, dict):
        result = {}
        for key, value in doc.items():
            result[key] = _serialize_value(value)
        return result
    if isinstance(doc, list):
        return [_serialize_value(item) for item in doc]
    return _serialize_value(doc)


def _serialize_value(value: Any) -> Any:
    """Serialize a value for JSON recording."""
    if value is None:
        return None
    if isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, bytes):
        return {"$binary": value.hex()}
    if isinstance(value, dict):
        return _serialize_document(value)
    if isinstance(value, list):
        return [_serialize_value(item) for item in value]
    # Handle ObjectId
    if hasattr(value, "__class__") and value.__class__.__name__ == "ObjectId":
        return {"$oid": str(value)}
    # Handle datetime
    if hasattr(value, "isoformat"):
        return {"$date": value.isoformat()}
    # Handle other BSON types
    return str(value)


def _hash_filter(filter_doc: Optional[Dict]) -> str:
    """Create a hash of the filter for matching."""
    if filter_doc is None:
        return hashlib.sha256(b"null").hexdigest()[:16]
    serialized = json.dumps(_serialize_document(filter_doc), sort_keys=True)
    return hashlib.sha256(serialized.encode()).hexdigest()[:16]


def _record_unsupported(collection_name: str, db_name: str, operation: str) -> None:
    """Record an unsupported operation event."""
    recorder = _get_recorder()
    recorder.record_event({
        "type": "MONGODB_UNSUPPORTED",
        "database": db_name,
        "collection": collection_name,
        "operation": operation,
        "message": f"MongoDB {operation} not supported for replay",
    })


class _RecordingCursor:
    """Wrapper around pymongo Cursor that records results."""

    def __init__(self, cursor, collection_name: str, db_name: str, operation: str,
                 filter_doc: Optional[Dict], pipeline: Optional[List] = None):
        self._cursor = cursor
        self._collection_name = collection_name
        self._db_name = db_name
        self._operation = operation
        self._filter_doc = filter_doc
        self._pipeline = pipeline
        self._results: List[Dict] = []
        self._recorded = False

    def __iter__(self):
        return self

    def __next__(self):
        try:
            doc = next(self._cursor)
            self._results.append(_serialize_document(doc))
            return doc
        except StopIteration:
            self._record_results()
            raise

    def _record_results(self) -> None:
        """Record the cursor results when iteration completes."""
        if self._recorded:
            return
        self._recorded = True

        recorder = _get_recorder()
        payload = {
            "database": self._db_name,
            "collection": self._collection_name,
            "operation": self._operation,
            "filter_hash": _hash_filter(self._filter_doc),
            "results": self._results,
            "result_count": len(self._results),
        }
        if self._pipeline:
            payload["pipeline_hash"] = _hash_filter({"pipeline": self._pipeline})
        recorder.record_mongodb(payload)

    def __getattr__(self, name):
        # Delegate other attributes to the underlying cursor
        return getattr(self._cursor, name)

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self._record_results()
        if hasattr(self._cursor, "__exit__"):
            return self._cursor.__exit__(*args)


def _patched_find(self, filter=None, *args, **kwargs):
    """Patched Collection.find that captures results."""
    cursor = _original_find(self, filter, *args, **kwargs)
    return _RecordingCursor(
        cursor,
        self.name,
        self.database.name,
        "find",
        filter,
    )


def _patched_find_one(self, filter=None, *args, **kwargs):
    """Patched Collection.find_one that captures result."""
    result = _original_find_one(self, filter, *args, **kwargs)

    recorder = _get_recorder()
    payload = {
        "database": self.database.name,
        "collection": self.name,
        "operation": "find_one",
        "filter_hash": _hash_filter(filter),
        "result": _serialize_document(result),
        "found": result is not None,
    }
    recorder.record_mongodb(payload)

    return result


def _patched_count_documents(self, filter, *args, **kwargs):
    """Patched Collection.count_documents that captures count."""
    count = _original_count_documents(self, filter, *args, **kwargs)

    recorder = _get_recorder()
    payload = {
        "database": self.database.name,
        "collection": self.name,
        "operation": "count_documents",
        "filter_hash": _hash_filter(filter),
        "count": count,
    }
    recorder.record_mongodb(payload)

    return count


def _patched_aggregate(self, pipeline, *args, **kwargs):
    """Patched Collection.aggregate that captures results."""
    cursor = _original_aggregate(self, pipeline, *args, **kwargs)
    return _RecordingCursor(
        cursor,
        self.name,
        self.database.name,
        "aggregate",
        None,
        pipeline=pipeline,
    )


def _patched_distinct(self, key, filter=None, *args, **kwargs):
    """Patched Collection.distinct that captures results."""
    result = _original_distinct(self, key, filter, *args, **kwargs)

    recorder = _get_recorder()
    payload = {
        "database": self.database.name,
        "collection": self.name,
        "operation": "distinct",
        "key": key,
        "filter_hash": _hash_filter(filter),
        "results": [_serialize_value(v) for v in result],
        "result_count": len(result),
    }
    recorder.record_mongodb(payload)

    return result


def _create_write_wrapper(operation: str, original_func: Callable) -> Callable:
    """Create a wrapper for write operations that records unsupported event."""
    def wrapper(self, *args, **kwargs):
        _record_unsupported(self.name, self.database.name, operation)
        return original_func(self, *args, **kwargs)
    return wrapper


def install() -> None:
    """Install pymongo capture patches."""
    global _original_find, _original_find_one, _original_count_documents
    global _original_aggregate, _original_distinct

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

    if _original_find is not None:
        return  # Already patched

    # 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

    # Wrap write operations to record unsupported events
    for op in UNSUPPORTED_WRITE_OPS:
        if hasattr(Collection, op):
            original = getattr(Collection, op)
            setattr(Collection, op, _create_write_wrapper(op, original))


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

    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


def get_pymongo_info() -> Dict[str, Any]:
    """Get pymongo detection info for manifest."""
    result = {"detected": False}

    try:
        import pymongo
        result["detected"] = True
        result["pymongo_version"] = pymongo.version
    except ImportError:
        pass

    if result["detected"]:
        result["supported_ops"] = list(SUPPORTED_READ_OPS)
        result["unsupported_ops"] = list(UNSUPPORTED_WRITE_OPS)

    return result
