"""
Metadata Store Client - Database operations for metadata storage.

Contract-centric API: all methods are keyed by (contract_name, contract_version).
Manages: schemas, coercion rules, validation rules, metadata records per contract.

Architecture notes:
    ``MetadataStoreClient`` is an abstract base class.  Subclasses (SQLite,
    Postgres, MongoDB, Redis, InMemory) implement the ``@abstractmethod``
    methods.  Template methods (``get_contract``, ``get_complete_schema``)
    are concrete and compose from the abstract primitives.

    The ``connect`` / ``disconnect`` / ``__enter__`` / ``__exit__`` lifecycle
    mirrors pystator's ``StateStoreClient``.
"""

from __future__ import annotations

import copy
from abc import ABC, abstractmethod
from typing import Any


class MetadataStoreClient(ABC):
    """
    Client for storing and retrieving metadata.  All operations are keyed by
    data contract (contract_name, contract_version).  Implementations extend
    this for specific backends (PostgreSQL, SQLite, MongoDB, etc.).
    """

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

        Args:
            connection_string: Database connection string (format depends on implementation)
        """
        self.connection_string = connection_string
        self._connection = None

    # -- lifecycle -----------------------------------------------------------

    @abstractmethod
    def connect(self) -> None:
        """Establish database connection.  Subclasses must implement."""

    def disconnect(self) -> None:
        """Close database connection."""
        if self._connection:
            self._connection = None

    # -- guard ---------------------------------------------------------------

    def _require_connection(self) -> None:
        """Raise ``RuntimeError`` if ``connect()`` has not been called."""
        if self._connection is None:
            raise RuntimeError("Not connected. Call connect() first.")

    # -------------------------------------------------------------------------
    # Contract-centric API  (abstract methods)
    # -------------------------------------------------------------------------

    @abstractmethod
    def store_schema(
        self,
        contract_name: str,
        contract_version: str,
        schema: dict[str, Any],
    ) -> str:
        """
        Store a JSON Schema for the given data contract.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.
            schema: JSON Schema dictionary (may contain "version" for schema artifact version).

        Returns:
            Contract ID or identifier (implementation-defined).
        """

    @abstractmethod
    def get_schema(self, contract_name: str, contract_version: str) -> dict[str, Any]:
        """
        Retrieve the schema for the given data contract.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.

        Returns:
            Schema dictionary with version included, or None if not found.
        """

    @abstractmethod
    def list_contracts(self) -> list[dict[str, Any]]:
        """
        List all stored data contracts.

        Returns:
            List of dicts with at least "name" and "version"; may include id, schema_id, etc.
        """

    @abstractmethod
    def store_metadata(
        self,
        contract_name: str,
        contract_version: str,
        metadata: dict[str, Any],
    ) -> str:
        """
        Store metadata for the given data contract.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.
            metadata: Metadata dictionary.

        Returns:
            Metadata record ID or identifier.
        """

    @abstractmethod
    def get_metadata(self, contract_name: str, contract_version: str) -> dict[str, Any]:
        """
        Retrieve metadata for the given data contract.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.

        Returns:
            Metadata dictionary or None if not found.
        """

    @abstractmethod
    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.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.
            coercion_rules: Dict mapping field names to coercion function names.

        Returns:
            Coercion rules ID or identifier.
        """

    @abstractmethod
    def get_coercion_rules(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any]:
        """
        Retrieve coercion rules for the given data contract.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.

        Returns:
            Coercion rules dictionary or None if not found.
        """

    @abstractmethod
    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.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.
            validation_rules: Dict mapping field names to validation configs.

        Returns:
            Validation rules ID or identifier.
        """

    @abstractmethod
    def get_validation_rules(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any]:
        """
        Retrieve validation rules for the given data contract.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.

        Returns:
            Validation rules dictionary or None if not found.
        """

    def get_contract(
        self,
        contract_name: str,
        contract_version: str,
        include_metadata: bool = True,
        include_ownership: bool = True,
        include_governance: bool = True,
    ) -> dict[str, Any]:
        """
        Return the full data contract (schema + coercion_rules + validation_rules + metadata)
        for the given contract name and version.

        Default implementation assembles from get_schema, get_coercion_rules,
        get_validation_rules, get_metadata. Subclasses may override for efficiency.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.
            include_metadata: Whether to include metadata in the contract.
            include_ownership: Whether to include ownership from metadata.
            include_governance: Whether to include governance_rules from metadata.

        Returns:
            Consolidated contract dict (schema, coercion_rules, validation_rules,
            metadata, versions) or None if contract/schema not found.
        """
        schema = self.get_schema(contract_name, contract_version)
        if not schema:
            return None
        if "version" not in schema:
            schema = dict(schema)
            schema["version"] = "1.0.0"
        coercion_rules = self.get_coercion_rules(contract_name, contract_version)
        validation_rules = self.get_validation_rules(contract_name, contract_version)
        metadata = None
        if include_metadata:
            metadata = self.get_metadata(contract_name, contract_version)
        ownership = None
        governance_rules = None
        if metadata:
            ownership = metadata.get("ownership") if include_ownership else None
            governance_rules = (
                metadata.get("governance_rules") if include_governance else None
            )
        # Build contract dict (same shape as build_contract from contract_builder)
        versions: dict[str, str | None] = {}
        if schema and "version" in schema:
            versions["schema"] = schema["version"]
        contract: dict[str, Any] = {
            "schema": copy.deepcopy(schema),
            "coercion_rules": copy.deepcopy(coercion_rules) if coercion_rules else {},
            "validation_rules": (
                copy.deepcopy(validation_rules) if validation_rules else {}
            ),
        }
        if versions:
            contract["versions"] = versions
        if include_metadata and metadata:
            meta_copy = copy.deepcopy(metadata)
            meta_copy.pop("version", None)
            if include_ownership and ownership:
                meta_copy["ownership"] = ownership
            if include_governance and governance_rules:
                meta_copy["governance_rules"] = governance_rules
            contract["metadata"] = meta_copy
        if include_ownership and ownership:
            contract["ownership"] = copy.deepcopy(ownership)
        if include_governance and governance_rules:
            contract["governance_rules"] = copy.deepcopy(governance_rules)
        # Include ontology if available
        ontology = self.get_ontology(contract_name, contract_version)
        if ontology:
            contract["ontology"] = copy.deepcopy(ontology)
        return contract

    # -- ontology (wiki) ---------------------------------------------------

    def store_ontology(
        self,
        contract_name: str,
        contract_version: str,
        ontology: dict[str, Any],
    ) -> str:
        """
        Store ontology annotations for the given data contract.

        Default implementation raises NotImplementedError. Override in
        store implementations that support ontology storage.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.
            ontology: Ontology dict (version, fields, etc.).

        Returns:
            Identifier for the stored ontology.
        """
        raise NotImplementedError(
            f"{type(self).__name__} does not support ontology storage. "
            "Use a wiki-enabled store implementation."
        )

    def get_ontology(self, contract_name: str, contract_version: str) -> dict[str, Any]:
        """
        Retrieve ontology annotations for the given data contract.

        Default implementation returns None (ontology not available).

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.

        Returns:
            Ontology dictionary or None if not found/supported.
        """
        return None

    # -- complete schema ----------------------------------------------------

    def get_complete_schema(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any]:
        """
        Retrieve the schema for the given contract with coercion and validation
        rules merged into the schema properties.

        Args:
            contract_name: Data contract name.
            contract_version: Data contract version.

        Returns:
            Complete schema dict with rules merged, or None if not found.
        """
        schema = self.get_schema(contract_name, contract_version)
        if not schema:
            return None
        if "version" not in schema:
            raise ValueError(
                f"Contract {contract_name!r}@{contract_version!r} schema has no version field."
            )
        complete_schema = copy.deepcopy(schema)
        coercion_rules = self.get_coercion_rules(contract_name, contract_version)
        if coercion_rules:
            _merge_coercion_rules(complete_schema, coercion_rules)
        validation_rules = self.get_validation_rules(contract_name, contract_version)
        if validation_rules:
            _merge_validation_rules(complete_schema, validation_rules)
        return complete_schema

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()


def _merge_coercion_rules(
    schema: dict[str, Any], coercion_rules: dict[str, Any] | None
) -> None:
    """Merge coercion rules into schema properties (in place)."""
    if "properties" not in schema:
        return
    for field_name, coercion_name in coercion_rules.items():
        if field_name in schema["properties"]:
            schema["properties"][field_name]["coercion"] = coercion_name


def _merge_validation_rules(
    schema: dict[str, Any], validation_rules: dict[str, Any] | None
) -> None:
    """Merge validation rules into schema properties (in place)."""
    if "properties" not in schema:
        return
    for field_path, field_validations in validation_rules.items():
        if "." in field_path:
            parts = field_path.split(".")
            if len(parts) == 2:
                parent_field, child_field = parts
                if parent_field in schema["properties"]:
                    parent_prop = schema["properties"][parent_field]
                    if (
                        "properties" in parent_prop
                        and child_field in parent_prop["properties"]
                    ):
                        child = parent_prop["properties"][child_field]
                        if "validations" not in child:
                            child["validations"] = {}
                        child["validations"].update(field_validations)
        else:
            if field_path in schema["properties"]:
                prop = schema["properties"][field_path]
                if "validations" not in prop:
                    prop["validations"] = {}
                prop["validations"].update(field_validations)
