"""
MongoDB Metadata Store Implementation

Stores metadata in MongoDB collections.
"""

from __future__ import annotations

from typing import Any

try:
    from pymongo import MongoClient  # type: ignore[import-untyped]
    from pymongo.collection import Collection  # type: ignore[import-untyped]
    from pymongo.database import Database  # type: ignore[import-untyped]

    MONGO_AVAILABLE = True
except ImportError:
    MONGO_AVAILABLE = False
    MongoClient = None  # type: ignore[assignment,misc]
    Collection = None  # type: ignore[assignment,misc]
    Database = None  # type: ignore[assignment,misc]

from pycharter.metadata_store.client import MetadataStoreClient


class MongoDBMetadataStore(MetadataStoreClient):
    """
    MongoDB metadata store implementation.

    Stores metadata in MongoDB collections:
    - data_contracts: Central collection that links all components together
    - schemas: JSON Schema definitions (linked via data_contract_id)
    - metadata_records: Metadata including governance rules (linked via data_contract_id)
    - coercion_rules: Coercion rules for data transformation (linked via data_contract_id)
    - validation_rules: Validation rules for data validation (linked via data_contract_id)

    Connection string format: mongodb://[username:password@]host[:port][/database]

    The structure mirrors PostgreSQL's data_contracts table, where all components
    are linked via data_contract_id. This ensures consistency across both storage backends.

    Example:
        >>> store = MongoDBMetadataStore("mongodb://localhost:27017/pycharter")
        >>> store.connect()  # Indexes are created automatically on first connection
        >>> store.store_schema("user", "1.0", {"type": "object"})
        >>> store.store_coercion_rules("user", "1.0", {"age": "coerce_to_integer"})
        >>> store.store_validation_rules("user", "1.0", {"age": {"is_positive": {}}})

    Note:
        Indexes are created automatically on connect() by default. The implementation
        checks if indexes exist before creating them to avoid unnecessary overhead.
        To skip index creation (e.g., if indexes are managed externally), use:
        >>> store.connect(ensure_indexes=False)
    """

    def __init__(
        self,
        connection_string: str | None = None,
        database_name: str = "pycharter",
    ):
        """
        Initialize MongoDB metadata store.

        Args:
            connection_string: MongoDB connection string
            database_name: Database name (default: "pycharter")
        """
        if not MONGO_AVAILABLE:
            raise ImportError(
                "pymongo is required for MongoDBMetadataStore. "
                "Install with: pip install pymongo"
            )
        if connection_string and not connection_string.startswith(
            ("mongodb://", "mongodb+srv://")
        ):
            raise ValueError(
                f"MongoDBMetadataStore requires a mongodb:// URL, got: {connection_string!r}"
            )
        super().__init__(connection_string)
        self.database_name = database_name
        self._client: MongoClient | None = None
        self._db: Any = None

    def connect(self, ensure_indexes: bool = True) -> None:
        """
        Connect to MongoDB.

        Args:
            ensure_indexes: If True, ensure indexes exist (default: True).
                          Set to False if you've already created indexes manually
                          or want to skip index creation for performance.
        """
        if not self.connection_string:
            raise ValueError("connection_string is required for MongoDB")

        self._client = MongoClient(self.connection_string)
        self._db = self._client[self.database_name]  # type: ignore[assignment]
        self._connection = self._db

        # Create indexes for better query performance (only if requested)
        if ensure_indexes:
            self._ensure_indexes()

    def disconnect(self) -> None:
        """Close MongoDB connection."""
        if self._client:
            self._client.close()
            self._client = None
            self._db = None
            self._connection = None

    def _ensure_indexes(self) -> None:
        """
        Ensure all required indexes exist.

        This method checks if indexes exist before creating them to avoid
        unnecessary overhead on every connection. Indexes are created in
        the background to avoid blocking operations.

        Note: MongoDB's create_index() is idempotent, but checking first
        avoids the overhead of the check that MongoDB does internally.
        """
        if self._db is None:
            return

        # Helper to normalize index spec for comparison
        def normalize_index_spec(spec):
            """Normalize index spec to a comparable format."""
            if isinstance(spec, str):
                return {spec: 1}
            elif isinstance(spec, list):
                return dict(spec)
            return spec

        # Helper to check if index exists by comparing key specs
        def index_exists_by_keys(collection, index_spec) -> bool:
            """Check if an index with the given key specification exists."""
            try:
                normalized_spec = normalize_index_spec(index_spec)
                indexes = list(collection.list_indexes())
                for idx in indexes:
                    # Compare index keys (excluding _id index)
                    idx_keys = idx.get("key", {})
                    if idx_keys == normalized_spec:
                        return True
                return False
            except Exception:
                # If we can't check, assume it doesn't exist and let create_index handle it
                return False

        # Helper to create index if it doesn't exist
        def create_index_if_missing(collection, index_spec, **kwargs):
            """Create index only if it doesn't exist."""
            if not index_exists_by_keys(collection, index_spec):
                collection.create_index(index_spec, background=True, **kwargs)

        # Create indexes for data_contracts collection
        create_index_if_missing(self._db.data_contracts, "name")
        create_index_if_missing(
            self._db.data_contracts, [("name", 1), ("version", 1)], unique=True
        )

        # Create indexes for schemas collection
        create_index_if_missing(self._db.schemas, "data_contract_id")
        create_index_if_missing(
            self._db.schemas, [("title", 1), ("version", 1)], unique=True
        )

        # Create indexes for metadata_records collection
        create_index_if_missing(self._db.metadata_records, "data_contract_id")
        create_index_if_missing(
            self._db.metadata_records,
            [("title", 1), ("version", 1)],
            unique=True,
        )

        # Create indexes for coercion_rules collection
        create_index_if_missing(self._db.coercion_rules, "data_contract_id")
        create_index_if_missing(self._db.coercion_rules, "schema_id")
        create_index_if_missing(
            self._db.coercion_rules,
            [("title", 1), ("version", 1)],
            unique=True,
        )

        # Create indexes for validation_rules collection
        create_index_if_missing(self._db.validation_rules, "data_contract_id")
        create_index_if_missing(self._db.validation_rules, "schema_id")
        create_index_if_missing(
            self._db.validation_rules,
            [("title", 1), ("version", 1)],
            unique=True,
        )

    def _get_or_create_data_contract(
        self,
        contract_name: str,
        version: str,
        status: str = "active",
        description: str | None = None,
    ) -> str:
        """Get or create a data_contract record. Returns data contract ID as string."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")
        from pycharter.shared.name_validator import validate_name

        contract_name = validate_name(contract_name, field_name="contract_name")

        existing = self._db.data_contracts.find_one(
            {"name": contract_name, "version": version}
        )
        if existing:
            return str(existing["_id"])

        from bson import ObjectId

        data_contract_id = ObjectId()
        doc = {
            "_id": data_contract_id,
            "name": contract_name,
            "version": version,
            "status": status,
            "description": description,
        }
        self._db.data_contracts.insert_one(doc)
        return str(data_contract_id)

    def _resolve_data_contract(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any] | None:
        """Resolve (contract_name, contract_version) to data_contract doc. Returns dict with id, schema_id, etc. or None."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")
        doc = self._db.data_contracts.find_one(
            {"name": contract_name, "version": contract_version}
        )
        if not doc:
            return None
        return {
            "id": doc["_id"],
            "schema_id": doc.get("schema_id"),
            "coercion_rules_id": doc.get("coercion_rules_id"),
            "validation_rules_id": doc.get("validation_rules_id"),
            "metadata_record_id": doc.get("metadata_record_id"),
        }

    def _get_schema_by_id(self, schema_id: Any) -> dict[str, Any] | None:
        """Load schema by _id (internal use)."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")
        doc = self._db.schemas.find_one({"_id": schema_id})
        if not doc:
            return None
        schema_data = doc.get("schema_data") or doc.get("schema")
        if not schema_data:
            return None
        if "version" not in schema_data:
            schema_data = dict(schema_data)
            schema_data["version"] = doc.get("version") or "1.0.0"
        return schema_data

    def store_schema(
        self,
        contract_name: str,
        contract_version: str,
        schema: dict[str, Any],
    ) -> str:
        """Store a schema for the given data contract. Returns contract ID (string)."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")

        schema_artifact_version = schema.get("version")
        if not schema_artifact_version:
            schema = dict(schema)
            schema["version"] = contract_version
            schema_artifact_version = contract_version
        elif "version" not in schema:
            schema = dict(schema)
            schema["version"] = schema_artifact_version

        from bson import ObjectId

        data_contract_id_str = self._get_or_create_data_contract(
            contract_name=contract_name,
            version=contract_version,
            description=schema.get("description"),
        )
        data_contract_id: Any = ObjectId(data_contract_id_str)

        from pycharter.shared.name_validator import (
            validate_name,
            validate_schema_table_names,
        )

        title = schema.get("title") or contract_name
        if title:
            title = validate_name(str(title), field_name="schema.title")

        validate_schema_table_names(schema)

        existing = self._db.schemas.find_one(
            {"title": title, "version": schema_artifact_version}
        )
        if existing:
            raise ValueError(
                f"Schema with title '{title}' and version '{schema_artifact_version}' already exists."
            )
        schema_id = ObjectId()
        doc = {
            "_id": schema_id,
            "title": title,
            "data_contract_id": data_contract_id,
            "version": schema_artifact_version,
            "schema_data": schema,
        }
        self._db.schemas.insert_one(doc)
        self._db.data_contracts.update_one(
            {"_id": data_contract_id},
            {"$set": {"schema_id": schema_id}},
        )
        return data_contract_id_str

    def get_schema(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any] | None:
        """Retrieve the schema for the given data contract."""
        row = self._resolve_data_contract(contract_name, contract_version)
        if not row or not row.get("schema_id"):
            return None
        return self._get_schema_by_id(row["schema_id"])

    def list_contracts(self) -> list[dict[str, Any]]:
        """List all data contracts with name, version, and optional id/schema_id."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")
        result = []
        for doc in self._db.data_contracts.find(
            {}, {"_id": 1, "name": 1, "version": 1, "schema_id": 1}
        ):
            result.append(
                {
                    "id": str(doc["_id"]),
                    "name": doc["name"],
                    "version": doc["version"],
                    "schema_id": (
                        str(doc["schema_id"]) if doc.get("schema_id") else None
                    ),
                }
            )
        return result

    def store_metadata(
        self,
        contract_name: str,
        contract_version: str,
        metadata: dict[str, Any],
    ) -> str:
        """Store metadata for the given data contract. Contract must exist."""
        row = self._resolve_data_contract(contract_name, contract_version)
        if not row:
            raise ValueError(
                f"Data contract not found: {contract_name!r} @ {contract_version!r}. Store the schema first."
            )
        data_contract_id = row["id"]

        from pycharter.shared.name_validator import normalize_name, validate_name

        title = metadata.get("title")
        if not title:
            title = (
                normalize_name(f"metadata_for_{contract_name}_{contract_version}")
                or "metadata"
            )
        normalized_title = normalize_name(str(title))
        if not normalized_title:
            raise ValueError(
                f"metadata.title '{title}' cannot be normalized to a valid name."
            )
        title = validate_name(normalized_title, field_name="metadata.title")
        version = contract_version
        status = metadata.get("status", "active")
        description = metadata.get("description")
        governance_rules = metadata.get("governance_rules")

        from bson import ObjectId

        existing = self._db.metadata_records.find_one(
            {"title": title, "version": version}
        )
        if existing:
            metadata_id = existing["_id"]
            self._db.metadata_records.update_one(
                {"_id": metadata_id},
                {
                    "$set": {
                        "title": title,
                        "status": status,
                        "description": description,
                        "governance_rules": governance_rules,
                    }
                },
            )
        else:
            metadata_id = ObjectId()
            self._db.metadata_records.insert_one(
                {
                    "_id": metadata_id,
                    "title": title,
                    "data_contract_id": data_contract_id,
                    "version": version,
                    "status": status,
                    "description": description,
                    "governance_rules": governance_rules,
                }
            )
        self._db.data_contracts.update_one(
            {"_id": data_contract_id},
            {"$set": {"metadata_record_id": metadata_id}},
        )
        return str(metadata_id)

    def get_metadata(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any] | None:
        """Retrieve metadata for the given data contract."""
        row = self._resolve_data_contract(contract_name, contract_version)
        if not row or not row.get("metadata_record_id"):
            return None
        return self._get_metadata_by_id(row["metadata_record_id"])

    def store_coercion_rules(
        self,
        contract_name: str,
        contract_version: str,
        coercion_rules: dict[str, Any],
    ) -> str:
        """Store coercion rules for the given data contract. Contract must have a schema."""
        row = self._resolve_data_contract(contract_name, contract_version)
        if not row:
            raise ValueError(
                f"Data contract not found: {contract_name!r} @ {contract_version!r}. Store the schema first."
            )
        data_contract_id = row["id"]
        schema_id = row.get("schema_id")
        if not schema_id:
            raise ValueError(
                f"Data contract {contract_name!r} @ {contract_version!r} has no schema. Store the schema first."
            )

        from bson import ObjectId

        from pycharter.shared.name_validator import normalize_name, validate_name

        title = (
            normalize_name(f"{contract_name}_coercion_rules")
            or normalize_name(f"{contract_name}_coercion")
            or "coercion_rules"
        )
        title = validate_name(title, field_name="coercion_rules.title")
        version = contract_version

        existing = self._db.coercion_rules.find_one(
            {"title": title, "version": version}
        )
        if existing:
            rule_id = existing["_id"]
            self._db.coercion_rules.update_one(
                {"_id": rule_id},
                {
                    "$set": {
                        "rules": coercion_rules,
                        "title": title,
                        "schema_id": schema_id,
                    }
                },
            )
        else:
            rule_id = ObjectId()
            self._db.coercion_rules.insert_one(
                {
                    "_id": rule_id,
                    "title": title,
                    "data_contract_id": data_contract_id,
                    "version": version,
                    "rules": coercion_rules,
                    "schema_id": schema_id,
                }
            )
        self._db.data_contracts.update_one(
            {"_id": data_contract_id},
            {"$set": {"coercion_rules_id": rule_id}},
        )
        return str(rule_id)

    def get_coercion_rules(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any] | None:
        """Retrieve coercion rules for the given data contract."""
        row = self._resolve_data_contract(contract_name, contract_version)
        if not row or not row.get("coercion_rules_id"):
            return None
        return self._get_coercion_rules_by_id(row["coercion_rules_id"])

    def store_validation_rules(
        self,
        contract_name: str,
        contract_version: str,
        validation_rules: dict[str, Any],
    ) -> str:
        """Store validation rules for the given data contract. Contract must have a schema."""
        row = self._resolve_data_contract(contract_name, contract_version)
        if not row:
            raise ValueError(
                f"Data contract not found: {contract_name!r} @ {contract_version!r}. Store the schema first."
            )
        data_contract_id = row["id"]
        schema_id = row.get("schema_id")
        if not schema_id:
            raise ValueError(
                f"Data contract {contract_name!r} @ {contract_version!r} has no schema. Store the schema first."
            )

        from bson import ObjectId

        from pycharter.shared.name_validator import normalize_name, validate_name

        title = (
            normalize_name(f"{contract_name}_validation_rules")
            or normalize_name(f"{contract_name}_validation")
            or "validation_rules"
        )
        title = validate_name(title, field_name="validation_rules.title")
        version = contract_version

        existing = self._db.validation_rules.find_one(
            {"title": title, "version": version}
        )
        if existing:
            rule_id = existing["_id"]
            self._db.validation_rules.update_one(
                {"_id": rule_id},
                {
                    "$set": {
                        "rules": validation_rules,
                        "title": title,
                        "schema_id": schema_id,
                    }
                },
            )
        else:
            rule_id = ObjectId()
            self._db.validation_rules.insert_one(
                {
                    "_id": rule_id,
                    "title": title,
                    "data_contract_id": data_contract_id,
                    "version": version,
                    "rules": validation_rules,
                    "schema_id": schema_id,
                }
            )
        self._db.data_contracts.update_one(
            {"_id": data_contract_id},
            {"$set": {"validation_rules_id": rule_id}},
        )
        return str(rule_id)

    def get_validation_rules(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any] | None:
        """Retrieve validation rules for the given data contract."""
        row = self._resolve_data_contract(contract_name, contract_version)
        if not row or not row.get("validation_rules_id"):
            return None
        return self._get_validation_rules_by_id(row["validation_rules_id"])

    def get_schema_info(self) -> dict[str, Any]:
        """
        Get information about the current database schema.

        Returns:
            Dictionary with schema information:
            {
                "revision": str or None,
                "initialized": bool,
                "message": str
            }
        """
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")

        # Check if collections exist (equivalent to "initialized")
        initialized = (
            "data_contracts" in self._db.list_collection_names()
            and "schemas" in self._db.list_collection_names()
        )

        # MongoDB doesn't have migrations like Postgres, so no revision
        revision = None

        return {
            "revision": revision,
            "initialized": initialized,
            "message": f"Schema initialized: {initialized}"
            + (f" (revision: {revision})" if revision else ""),
        }

    def _get_coercion_rules_by_id(
        self, coercion_rules_id: Any
    ) -> dict[str, Any] | None:
        """Load coercion rules by _id (internal use)."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")

        from bson import ObjectId

        try:
            rule_obj_id: Any = ObjectId(coercion_rules_id)
        except Exception:
            rule_obj_id = coercion_rules_id

        doc = self._db.coercion_rules.find_one({"_id": rule_obj_id})
        if doc:
            return doc.get("rules")
        return None

    def _get_validation_rules_by_id(
        self, validation_rules_id: Any
    ) -> dict[str, Any] | None:
        """Load validation rules by _id (internal use)."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")

        from bson import ObjectId

        try:
            rule_obj_id: Any = ObjectId(validation_rules_id)
        except Exception:
            rule_obj_id = validation_rules_id

        doc = self._db.validation_rules.find_one({"_id": rule_obj_id})
        if doc:
            return doc.get("rules")
        return None

    def _get_metadata_by_id(self, metadata_id: Any) -> dict[str, Any] | None:
        """Load metadata by _id (internal use)."""
        if self._db is None:
            raise RuntimeError("Not connected. Call connect() first.")

        from bson import ObjectId

        try:
            metadata_obj_id: Any = ObjectId(metadata_id)
        except Exception:
            metadata_obj_id = metadata_id

        doc = self._db.metadata_records.find_one({"_id": metadata_obj_id})
        if doc:
            # Reconstruct metadata dictionary
            metadata = {
                "title": doc.get("title"),
                "status": doc.get("status"),
                "type": doc.get("type"),
                "description": doc.get("description"),
                "version": doc.get("version"),
            }

            # Add JSON fields
            if doc.get("governance_rules"):
                metadata["governance_rules"] = doc.get("governance_rules")

            # Add relationship fields if present
            relationship_fields = [
                "business_owners",
                "bu_sme",
                "it_application_owners",
                "it_sme",
                "support_lead",
                "domains",
                "system_sources",
                "system_pulls",
                "system_pushes",
            ]
            for field in relationship_fields:
                if field in doc:
                    metadata[field] = doc[field]

            return metadata
        return None
