"""Private shared SQLAlchemy base for SQLite and PostgreSQL metadata stores.

Provides all ORM query logic via ``DataContractModel``, ``SchemaModel``,
``CoercionRuleModel``, ``ValidationRuleModel``, and ``MetadataRecordModel``.
Public subclasses (``SQLiteMetadataStore``, ``PostgresMetadataStore``) only
validate connection strings and call ``_connect_engine`` with the correct
schema name.

This module is internal — import the public stores from
``pycharter.metadata_store``.
"""

from __future__ import annotations

import json
import logging
import uuid
from typing import Any

from sqlalchemy import inspect
from sqlalchemy.orm import Session, sessionmaker

from pycharter.db.models.base import Base, get_engine
from pycharter.db.models.coercion_rule import CoercionRuleModel
from pycharter.db.models.data_contract import DataContractModel
from pycharter.db.models.metadata_record import MetadataRecordModel
from pycharter.db.models.schema import SchemaModel
from pycharter.db.models.validation_rule import ValidationRuleModel
from pycharter.metadata_store.client import MetadataStoreClient

logger = logging.getLogger(__name__)


class _SQLAlchemyMetadataStoreBase(MetadataStoreClient):
    """Shared SQLAlchemy query logic for SQLite and PostgreSQL stores.

    Subclasses must:
    1. Validate the connection string in ``__init__``.
    2. Implement ``connect()`` and call ``_connect_engine()`` with the
       appropriate ``schema_name``.
    """

    def __init__(self, connection_string: str) -> None:
        super().__init__(connection_string)
        self._engine: Any = None
        self._session: Session | None = None
        self._schema_name: str | None = None

    # ------------------------------------------------------------------
    # Engine / session lifecycle
    # ------------------------------------------------------------------

    def _connect_engine(
        self,
        schema_name: str | None,
        *,
        validate_schema: bool = True,
        auto_initialize: bool = False,
    ) -> None:
        """Create engine, session, and optionally validate / create tables.

        Args:
            schema_name: Database schema to inspect (``None`` for SQLite,
                ``"pycharter"`` for PostgreSQL).
            validate_schema: Check that required tables exist.
            auto_initialize: Create tables if they are missing.
        """
        self._schema_name = schema_name
        self._engine = get_engine(self.connection_string)  # type: ignore[arg-type]
        factory = sessionmaker(bind=self._engine)
        self._session = factory()
        self._connection = self._session

        if validate_schema and not self._is_schema_initialized():
            if auto_initialize:
                self._initialize_schema()
            else:
                raise RuntimeError(
                    "Database schema is not initialized. "
                    "Please run 'pycharter db init' to initialize the schema first.\n"
                    f"Example: pycharter db init {self.connection_string}"
                )

    def disconnect(self) -> None:
        """Close session and dispose engine."""
        if self._session is not None:
            self._session.close()
            self._session = None
        if self._engine is not None:
            self._engine.dispose()
            self._engine = None
        self._connection = None

    # ------------------------------------------------------------------
    # Schema helpers
    # ------------------------------------------------------------------

    def _is_schema_initialized(self) -> bool:
        """Check whether the ``schemas`` table exists."""
        if self._engine is None:
            return False
        try:
            insp = inspect(self._engine)
            tables = insp.get_table_names(schema=self._schema_name)
            return "schemas" in tables
        except Exception:
            return False

    #: Tables required by the metadata store.  Avoids creating wiki tables
    #: that use PostgreSQL-specific types (JSONB) unsupported on SQLite.
    _CORE_TABLES = [
        DataContractModel.__table__,
        SchemaModel.__table__,
        MetadataRecordModel.__table__,
        CoercionRuleModel.__table__,
        ValidationRuleModel.__table__,
    ]

    def _initialize_schema(self) -> None:
        """Create core metadata tables from SQLAlchemy models."""
        try:
            Base.metadata.create_all(self._engine, tables=self._CORE_TABLES)
        except Exception as exc:
            raise RuntimeError(f"Failed to initialize schema: {exc}") from exc

    def get_schema_info(self) -> dict[str, Any]:
        """Get information about the current database schema."""
        self._require_connection()
        initialized = self._is_schema_initialized()
        revision: str | None = None

        if initialized:
            try:
                insp = inspect(self._engine)
                tables = insp.get_table_names(schema=self._schema_name)
                if "alembic_version" in tables:
                    from sqlalchemy import text

                    result = self._session.execute(  # type: ignore[union-attr]
                        text("SELECT version_num FROM alembic_version LIMIT 1")
                    )
                    row = result.fetchone()
                    if row:
                        revision = row[0]
            except Exception:
                pass

        message = f"Schema initialized: {initialized}"
        if revision:
            message += f" (revision: {revision})"

        return {
            "revision": revision,
            "initialized": initialized,
            "message": message,
        }

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _get_session(self) -> Session:
        """Return the active session, raising if not connected."""
        self._require_connection()
        assert self._session is not None  # noqa: S101
        return self._session

    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. Returns data contract ID string."""
        from pycharter.shared.name_validator import validate_name

        contract_name = validate_name(contract_name, field_name="contract_name")
        session = self._get_session()
        row = (
            session.query(DataContractModel)
            .filter(
                DataContractModel.name == contract_name,
                DataContractModel.version == version,
            )
            .first()
        )
        if row:
            return str(row.id)
        new_dc = DataContractModel(
            name=contract_name,
            version=version,
            status=status,
            description=description,
        )
        session.add(new_dc)
        session.flush()
        return str(new_dc.id)

    def _resolve_data_contract(
        self, contract_name: str, contract_version: str
    ) -> dict[str, Any] | None:
        """Resolve (contract_name, contract_version) to data_contract row."""
        session = self._get_session()
        row = (
            session.query(DataContractModel)
            .filter(
                DataContractModel.name == contract_name,
                DataContractModel.version == contract_version,
            )
            .first()
        )
        if not row:
            return None
        return {
            "id": str(row.id),
            "schema_id": str(row.schema_id) if row.schema_id else None,
            "coercion_rules_id": (
                str(row.coercion_rules_id) if row.coercion_rules_id else None
            ),
            "validation_rules_id": (
                str(row.validation_rules_id) if row.validation_rules_id else None
            ),
            "metadata_record_id": (
                str(row.metadata_record_id) if row.metadata_record_id else None
            ),
        }

    # ------------------------------------------------------------------
    # Schema operations
    # ------------------------------------------------------------------

    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."""
        session = self._get_session()

        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

        data_contract_id = self._get_or_create_data_contract(
            contract_name=contract_name,
            version=contract_version,
            description=schema.get("description"),
        )

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

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

        validate_schema_table_names(schema)

        # Check for duplicate
        existing = (
            session.query(SchemaModel)
            .filter(
                SchemaModel.title == title,
                SchemaModel.version == schema_artifact_version,
            )
            .first()
        )
        if existing:
            raise ValueError(
                f"Schema with title '{title}' and version "
                f"'{schema_artifact_version}' already exists."
            )

        schema_obj = SchemaModel(
            title=title,
            data_contract_id=uuid.UUID(data_contract_id),
            version=schema_artifact_version,
            schema_data=schema,
        )
        session.add(schema_obj)
        session.flush()

        # Update data_contract.schema_id
        dc = session.query(DataContractModel).get(uuid.UUID(data_contract_id))
        if dc:
            dc.schema_id = schema_obj.id
        session.commit()
        return data_contract_id

    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 _get_schema_by_id(self, schema_id: str) -> dict[str, Any] | None:
        """Load schema by id."""
        session = self._get_session()
        row = session.query(SchemaModel).get(uuid.UUID(schema_id))
        if not row:
            return None
        schema_data = row.schema_data
        if isinstance(schema_data, str):
            schema_data = json.loads(schema_data)
        if schema_data and "version" not in schema_data:
            schema_data = dict(schema_data)
            schema_data["version"] = row.version or "1.0.0"
        return schema_data

    def list_contracts(self) -> list[dict[str, Any]]:
        """List all data contracts."""
        session = self._get_session()
        rows = (
            session.query(DataContractModel)
            .order_by(DataContractModel.name, DataContractModel.version)
            .all()
        )
        return [
            {
                "id": str(r.id),
                "name": r.name,
                "version": r.version,
                "schema_id": str(r.schema_id) if r.schema_id else None,
            }
            for r in rows
        ]

    # ------------------------------------------------------------------
    # Metadata operations
    # ------------------------------------------------------------------

    def store_metadata(
        self,
        contract_name: str,
        contract_version: str,
        metadata: dict[str, Any],
    ) -> str:
        """Store metadata for the given data contract."""
        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"]
        session = self._get_session()

        from pycharter.db.metadata_normalizer import ensure_canonical
        from pycharter.shared.name_validator import normalize_name, validate_name

        canonical = ensure_canonical(metadata)
        title = canonical.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

        # Check for duplicate
        existing = (
            session.query(MetadataRecordModel)
            .filter(
                MetadataRecordModel.title == title,
                MetadataRecordModel.version == version,
            )
            .first()
        )
        if existing:
            raise ValueError(
                f"Metadata with title '{title}' and version '{version}' already exists."
            )

        meta_obj = MetadataRecordModel(
            title=title,
            data_contract_id=uuid.UUID(data_contract_id),
            version=version,
            status=canonical.get("status", "active"),
            type=canonical.get("type"),
            description=canonical.get("description"),
            governance_rules=canonical.get("governance_rules"),
            payload=canonical if canonical else None,
        )
        session.add(meta_obj)
        session.flush()

        # Update data_contract.metadata_record_id
        dc = session.query(DataContractModel).get(uuid.UUID(data_contract_id))
        if dc:
            dc.metadata_record_id = meta_obj.id
        session.commit()
        return str(meta_obj.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 _get_metadata_by_id(self, metadata_id: str) -> dict[str, Any] | None:
        """Load metadata by id."""
        session = self._get_session()
        row = session.query(MetadataRecordModel).get(uuid.UUID(metadata_id))
        if not row:
            return None
        # Prefer payload if available.
        if row.payload:
            payload = row.payload
            if isinstance(payload, str):
                try:
                    data = json.loads(payload)
                    if isinstance(data, dict) and data:
                        return data
                except (json.JSONDecodeError, TypeError):
                    pass
            elif isinstance(payload, dict) and payload:
                return payload
        # Fallback: reconstruct from columns.
        out: dict[str, Any] = {
            "title": row.title,
            "version": row.version,
            "status": row.status,
            "type": row.type,
            "description": row.description,
            "ownership": {},
            "lineage": {},
        }
        if row.governance_rules:
            gov = row.governance_rules
            if isinstance(gov, str):
                try:
                    out["governance_rules"] = json.loads(gov)
                except (json.JSONDecodeError, TypeError):
                    pass
            else:
                out["governance_rules"] = gov
        return out

    # ------------------------------------------------------------------
    # Coercion rules
    # ------------------------------------------------------------------

    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."""
        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."
            )
        session = self._get_session()

        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

        # Check for duplicate
        existing = (
            session.query(CoercionRuleModel)
            .filter(
                CoercionRuleModel.title == title,
                CoercionRuleModel.version == version,
            )
            .first()
        )
        if existing:
            raise ValueError(
                f"Coercion rules with title '{title}' and version '{version}' "
                "already exist."
            )

        rule_obj = CoercionRuleModel(
            title=title,
            data_contract_id=uuid.UUID(data_contract_id),
            version=version,
            rules=coercion_rules,
            schema_id=uuid.UUID(schema_id),
        )
        session.add(rule_obj)
        session.flush()

        dc = session.query(DataContractModel).get(uuid.UUID(data_contract_id))
        if dc:
            dc.coercion_rules_id = rule_obj.id
        session.commit()
        return str(rule_obj.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
        session = self._get_session()
        obj = session.query(CoercionRuleModel).get(uuid.UUID(row["coercion_rules_id"]))
        if not obj:
            return None
        rules = obj.rules
        if isinstance(rules, str):
            return json.loads(rules)
        return rules

    # ------------------------------------------------------------------
    # Validation rules
    # ------------------------------------------------------------------

    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."""
        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."
            )
        session = self._get_session()

        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

        # Check for duplicate
        existing = (
            session.query(ValidationRuleModel)
            .filter(
                ValidationRuleModel.title == title,
                ValidationRuleModel.version == version,
            )
            .first()
        )
        if existing:
            raise ValueError(
                f"Validation rules with title '{title}' and version '{version}' "
                "already exist."
            )

        rule_obj = ValidationRuleModel(
            title=title,
            data_contract_id=uuid.UUID(data_contract_id),
            version=version,
            rules=validation_rules,
            schema_id=uuid.UUID(schema_id),
        )
        session.add(rule_obj)
        session.flush()

        dc = session.query(DataContractModel).get(uuid.UUID(data_contract_id))
        if dc:
            dc.validation_rules_id = rule_obj.id
        session.commit()
        return str(rule_obj.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
        session = self._get_session()
        obj = session.query(ValidationRuleModel).get(
            uuid.UUID(row["validation_rules_id"])
        )
        if not obj:
            return None
        rules = obj.rules
        if isinstance(rules, str):
            return json.loads(rules)
        return rules
