from __future__ import annotations

import logging
import re
from typing import Any, Dict, List, Optional, Tuple

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.engine import Engine

from mnemon.db_connect import get_pool_config
from mnemon.utils.error_handling import normalize_error, to_api_response

from actor.observability import ComponentStatus, ObservabilityService

logger = logging.getLogger(__name__)

_SAFE_IDENTIFIER = re.compile(r"^[a-zA-Z0-9_.]+$")


def _make_serializable(obj: Any) -> Any:
    """Recursively convert non-JSON-serializable values to strings.

    Handles SQLAlchemy type objects (INTEGER, VARCHAR, etc.), Decimal,
    datetime, date, time, bytes, memoryview, and UUID.
    """
    if obj is None or isinstance(obj, (bool, int, float, str)):
        return obj
    if isinstance(obj, dict):
        return {k: _make_serializable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_make_serializable(v) for v in obj]
    # Decimal -> float
    import decimal
    if isinstance(obj, decimal.Decimal):
        return float(obj)
    # datetime/date/time -> ISO string
    import datetime as _dt
    if isinstance(obj, (_dt.datetime, _dt.date, _dt.time)):
        return obj.isoformat()
    # bytes/memoryview -> hex string
    if isinstance(obj, (bytes, memoryview)):
        return bytes(obj).hex()
    # UUID
    import uuid
    if isinstance(obj, uuid.UUID):
        return str(obj)
    # Anything else (SQLAlchemy types, etc.) -> str
    return str(obj)

try:
    from mnemon.utils.dialect_queries import DIALECT_QUERIES
except ImportError:  # pragma: no cover
    DIALECT_QUERIES = {}


def _extract_database_type(url: str) -> Optional[str]:
    """Extract database type from SQLAlchemy URL dialect.

    Examples:
        postgresql+psycopg2://... -> postgresql
        mysql+pymysql://... -> mysql
        snowflake://... -> snowflake
        bigquery://... -> bigquery
    """
    if not url:
        return None
    # Match dialect before optional +driver and ://
    match = re.match(r"^([a-zA-Z0-9_]+)(?:\+[a-zA-Z0-9_]+)?://", url)
    if match:
        dialect = match.group(1).lower()
        # Normalize some common variants
        dialect_map = {
            "mssql": "sqlserver",
            "awsathena": "athena",
            "redshift": "redshift",
            "oracle": "oracle",
            "teradatasql": "teradata",
            "db2": "db2",
            "hive": "sparksql",
            "vertica": "vertica",
            "clickhouse": "clickhouse",
        }
        return dialect_map.get(dialect, dialect)
    return None


def _validate_identifier(name: str) -> None:
    """Validate a SQL identifier (schema/table name) against injection."""
    if not _SAFE_IDENTIFIER.match(name):
        raise ValueError(f"Invalid SQL identifier: {name!r}")


class DBClient:
    """SQLAlchemy-based database client for validation, metadata, and read-only queries."""

    def __init__(
        self,
        url: str,
        connect_args: Optional[Dict[str, Any]],
        observability: ObservabilityService,
        database_type: Optional[str] = None,
    ):
        self.url = url
        self.connect_args = connect_args or {}
        self.observability = observability
        self.engine: Optional[Engine] = None
        self._database_type = database_type or _extract_database_type(url)
        self.observability.register_component("database")

    @property
    def database_type(self) -> Optional[str]:
        """Get the database type."""
        return self._database_type

    def _ensure_engine(self) -> Engine:
        """Create and cache SQLAlchemy engine with appropriate pool configuration."""
        if self.engine is None:
            # Get pool config based on database type
            pool_config = {}
            if self._database_type:
                try:
                    pool_config = get_pool_config(self._database_type)
                except Exception as exc:
                    logger.warning(f"Failed to get pool config for {self._database_type}: {exc}")

            self.engine = create_engine(
                self.url,
                connect_args=self.connect_args,
                **pool_config,
            )
            logger.info(f"Database engine created for {self._database_type or 'unknown'}")
        return self.engine

    def _get_dialect_query(self, query_name: str) -> Optional[str]:
        """Get a dialect-specific query template if available."""
        if self._database_type and self._database_type in DIALECT_QUERIES:
            queries = DIALECT_QUERIES[self._database_type]
            return queries.get(query_name)
        return None

    def validate(self) -> None:
        """Validate database connectivity by executing a test query."""
        engine = self._ensure_engine()
        try:
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            self.observability.update_component_health(
                "database",
                status=ComponentStatus.HEALTHY,
                details={"connected": True, "database_type": self._database_type},
            )
        except Exception as exc:
            self.observability.update_component_health(
                "database",
                status=ComponentStatus.UNHEALTHY,
                details={"connected": False, "database_type": self._database_type},
                error=str(exc),
            )
            raise

    def safe_connect(self) -> Tuple[bool, Optional[Engine], Optional[Dict[str, Any]]]:
        """Attempt connection with normalized error handling.

        Returns:
            Tuple of (success, engine_or_none, error_response_or_none)
        """
        try:
            engine = self._ensure_engine()
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            self.observability.update_component_health(
                "database",
                status=ComponentStatus.HEALTHY,
                details={"connected": True, "database_type": self._database_type},
            )
            return (True, engine, None)
        except Exception as exc:
            self.observability.update_component_health(
                "database",
                status=ComponentStatus.UNHEALTHY,
                details={"connected": False, "database_type": self._database_type},
                error=str(exc),
            )
            # Normalize the error for API response
            db_type = self._database_type or "unknown"
            normalized = normalize_error(db_type, exc)
            return (False, None, to_api_response(normalized))

    def list_schemas(self) -> List[str]:
        """List all schema names in the database."""
        engine = self._ensure_engine()
        query_tpl = self._get_dialect_query("get_schemas")
        if query_tpl:
            try:
                with engine.connect() as conn:
                    result = conn.execute(text(query_tpl))
                    return [_make_serializable(row[0]) for row in result.fetchall()]
            except Exception:
                logger.debug("Dialect query for schemas failed; falling back to inspector", exc_info=True)
        insp = inspect(engine)
        return [_make_serializable(s) for s in insp.get_schema_names()]

    def list_tables(self, schema: Optional[str] = None) -> List[str]:
        """List all table names in the specified schema."""
        engine = self._ensure_engine()
        query_tpl = self._get_dialect_query("get_schema_tables")
        if query_tpl and schema:
            _validate_identifier(schema)
            try:
                with engine.connect() as conn:
                    result = conn.execute(text(query_tpl), {"schema": schema})
                    return [_make_serializable(row[0]) for row in result.fetchall()]
            except Exception:
                logger.debug("Dialect query for tables failed; falling back to inspector", exc_info=True)
        insp = inspect(engine)
        return [_make_serializable(t) for t in insp.get_table_names(schema=schema)]

    def describe_table(self, table: str, schema: Optional[str] = None) -> List[Dict[str, Any]]:
        """Get column information for a table."""
        engine = self._ensure_engine()
        query_tpl = self._get_dialect_query("get_table_columns")
        if query_tpl:
            _validate_identifier(table)
            if schema:
                _validate_identifier(schema)
            try:
                qualified = f"{schema}.{table}" if schema else table
                with engine.connect() as conn:
                    result = conn.execute(text(query_tpl.format(table=qualified)))
                    keys = result.keys()
                    return [_make_serializable(dict(zip(keys, row))) for row in result.fetchall()]
            except Exception:
                logger.debug("Dialect query for columns failed; falling back to inspector", exc_info=True)
        insp = inspect(engine)
        return [_make_serializable(col) for col in insp.get_columns(table, schema=schema)]

    def sample_table(self, table: str, schema: Optional[str] = None, limit: int = 5) -> List[Dict[str, Any]]:
        """Get sample rows from a table using dialect query or fallback."""
        _validate_identifier(table)
        if schema:
            _validate_identifier(schema)
        engine = self._ensure_engine()
        qualified = f"{schema}.{table}" if schema else table
        query_tpl = self._get_dialect_query("get_table_sample")
        if query_tpl:
            try:
                query_str = query_tpl.format(table=qualified, limit=int(limit))
                with engine.connect() as conn:
                    result = conn.execute(text(query_str))
                    keys = result.keys()
                    return [_make_serializable(dict(zip(keys, row))) for row in result.fetchall()]
            except Exception:
                logger.debug("Dialect query for sample failed; falling back", exc_info=True)
        # Fallback: generic SELECT
        query_str = f"SELECT * FROM {qualified} LIMIT {int(limit)}"
        with engine.connect() as conn:
            result = conn.execute(text(query_str))
            keys = result.keys()
            return [_make_serializable(dict(zip(keys, row))) for row in result.fetchall()]

    def table_count(self, table: str, schema: Optional[str] = None) -> int:
        """Get row count for a table using dialect query or fallback."""
        _validate_identifier(table)
        if schema:
            _validate_identifier(schema)
        engine = self._ensure_engine()
        qualified = f"{schema}.{table}" if schema else table
        query_tpl = self._get_dialect_query("get_table_count")
        if query_tpl:
            try:
                query_str = query_tpl.format(table=qualified)
                with engine.connect() as conn:
                    result = conn.execute(text(query_str))
                    row = result.fetchone()
                    return int(row[0]) if row else 0
            except Exception:
                logger.debug("Dialect query for count failed; falling back", exc_info=True)
        # Fallback: generic COUNT
        query_str = f"SELECT COUNT(*) FROM {qualified}"
        with engine.connect() as conn:
            result = conn.execute(text(query_str))
            row = result.fetchone()
            return int(row[0]) if row else 0

    def run_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """Execute a read-only query and return results as list of dicts."""
        self._assert_read_only(query)
        engine = self._ensure_engine()
        with engine.connect() as conn:
            result = conn.execute(text(query), params or {})
            if result.returns_rows:
                keys = result.keys()
                return [_make_serializable(dict(zip(keys, row))) for row in result.fetchall()]
            return []

    def safe_query(
        self, query: str, params: Optional[Dict[str, Any]] = None
    ) -> Tuple[bool, List[Dict[str, Any]], Optional[Dict[str, Any]]]:
        """Execute a query with normalized error handling.

        Returns:
            Tuple of (success, results, error_response_or_none)
        """
        try:
            results = self.run_query(query, params)
            return (True, results, None)
        except ValueError as exc:
            # Read-only validation failure
            return (False, [], {"success": False, "error": {"category": "invalid_query", "message": str(exc)}})
        except Exception as exc:
            db_type = self._database_type or "unknown"
            normalized = normalize_error(db_type, exc)
            return (False, [], to_api_response(normalized))

    def normalize_columns(self, columns: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Normalize SQLAlchemy column metadata to a stable schema."""
        normalized: List[Dict[str, Any]] = []
        for col in columns:
            if not isinstance(col, dict):
                continue
            name = col.get("name") or col.get("column_name")
            dtype = col.get("type") or col.get("data_type")
            is_nullable = col.get("nullable")
            if is_nullable is None:
                is_nullable = col.get("is_nullable")
            if isinstance(is_nullable, bool):
                is_nullable = "YES" if is_nullable else "NO"
            default_val = col.get("default") or col.get("column_default")
            normalized.append(
                {
                    "column_name": name,
                    "data_type": str(dtype) if dtype is not None else None,
                    "is_nullable": is_nullable,
                    "column_default": str(default_val) if default_val is not None else None,
                }
            )
        return normalized

    def _assert_read_only(self, query: str) -> None:
        """Validate that a query is read-only."""
        normalized = query.strip()
        # Strip block comments (/* ... */) first to prevent bypass
        normalized = re.sub(r'/\*.*?\*/', ' ', normalized, flags=re.DOTALL)
        # Strip simple inline comments
        lines = []
        for line in normalized.splitlines():
            for marker in ("--", "#"):
                idx = line.find(marker)
                if idx != -1:
                    line = line[:idx]
            lines.append(line)
        normalized = " ".join(lines).strip()
        lower = normalized.lower()
        allowed_starts = ("select", "with", "show", "describe", "explain", "pragma")
        first_token = lower.split()[0] if lower else ""
        if first_token not in allowed_starts:
            raise ValueError("Only read-only queries are permitted")
        disallowed = ("insert", "update", "delete", "alter", "drop", "truncate", "create", "replace", "grant", "revoke")
        if any(dml in lower for dml in disallowed):
            raise ValueError("Only read-only queries are permitted")

    def dispose(self) -> None:
        """Dispose of the database engine and release connections."""
        if self.engine:
            self.engine.dispose()
            self.engine = None
            logger.info("Database engine disposed")

    def __enter__(self) -> "DBClient":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.dispose()
