"""SQL database source helpers"""

import warnings
from typing import (
    Callable,
    Any,
    Dict,
    List,
    Literal,
    Optional,
    Iterator,
    Type,
    Union,
)
import operator
from abc import ABC, abstractmethod

import dlt
from dlt.common import logger
from dlt.common.configuration.specs import (
    BaseConfiguration,
    ConnectionStringCredentials,
    configspec,
)
from dlt.common.exceptions import (
    DltException,
    MissingDependencyException,
    ValueErrorWithKnownValues,
)
from dlt.common.schema import TTableSchemaColumns
from dlt.common.schema.typing import TWriteDispositionDict
from dlt.common.schema.utils import merge_columns
from dlt.common.typing import TColumnNames, TDataItem, TSortOrder, add_value_to_literal
from dlt.common.jsonpath import extract_simple_field_name
from dlt.common.utils import is_typeerror_due_to_wrong_call

from dlt.extract import Incremental
from dlt.extract.items_transform import LimitItem

from .arrow_helpers import row_tuples_to_arrow
from .schema_types import (
    default_table_adapter,
    Table,
    SelectAny,
    ReflectionLevel,
    TTypeAdapter,
    table_to_resource_hints,
)

from dlt.common.libs.sql_alchemy import (
    Engine,
    CompileError,
    create_engine,
    MetaData,
    sa,
    TextClause,
    ORACLE_NUMBER,
    OracleDialect,
)

TableBackend = Literal["sqlalchemy", "pyarrow", "pandas", "connectorx"]
SelectClause = Union[SelectAny, TextClause]
TQueryAdapter = Union[
    Callable[[SelectAny, Table], SelectClause],
    Callable[[SelectAny, Table, Incremental[Any], Engine], SelectClause],
]
TTableAdapter = Callable[[Table], Optional[Union[SelectAny, Table]]]

TABLE_LOADER_REGISTRY: Dict[str, Type["BaseTableLoader"]] = {}
"""Maps backend name to the table loader class that handles it."""


def _maybe_fix_0000_timezone(df: Any) -> Any:
    """Optionally convert +00:00 timezone to UTC on Arrow tables."""
    try:
        from dlt.common.libs.pyarrow import set_plus0000_timezone_to_utc, pyarrow

        # TODO: skip when Arrow releases timezone fix
        if isinstance(df, pyarrow.Table):
            return set_plus0000_timezone_to_utc(df)
    except MissingDependencyException:
        pass
    return df


class BaseTableLoader(ABC):
    """Base class for SQL table loaders.

    Provides query building infrastructure including incremental loading support.
    Subclasses must implement ``load_rows`` to handle the actual data retrieval.
    For loaders that use a SQLAlchemy connection, subclass ``TableLoader`` instead
    to reuse the result-conversion helpers.

    Args:
        engine: SQLAlchemy ``Engine`` used for query compilation and (optionally)
            execution.
        backend: Requested output format (``"sqlalchemy"``, ``"pyarrow"``,
            ``"pandas"``, ``"connectorx"``).
        table: Reflected SQLAlchemy ``Table`` object.
        columns: dlt column schema hints for the table.
        chunk_size: Number of rows per batch.
        incremental: Optional incremental loading state.
        query_adapter_callback: Optional callback to modify the generated query.
        limit: Optional row/time limit from the resource.
    """

    def __init__(
        self,
        engine: Engine,
        backend: TableBackend,
        table: Table,
        columns: TTableSchemaColumns,
        chunk_size: int = 1000,
        incremental: Optional[Incremental[Any]] = None,
        query_adapter_callback: Optional[TQueryAdapter] = None,
        limit: Optional[LimitItem] = None,
    ) -> None:
        self.engine = engine
        self.backend = backend
        self.table = table
        self.columns = columns
        self.chunk_size = chunk_size
        self.query_adapter_callback = query_adapter_callback
        self.incremental = incremental
        self.limit = limit
        if incremental:
            column_name = extract_simple_field_name(incremental.cursor_path)

            if column_name is None:
                raise ValueError(
                    f"Cursor path `{incremental.cursor_path}` must be a simple column name (e.g."
                    " `created_at`)"
                )

            try:
                self.cursor_column = table.c[column_name]
            except KeyError as e:
                raise KeyError(
                    f"Cursor column `{incremental.cursor_path}` does not exist in table"
                    f" `{table.name}`"
                ) from e
            self.last_value = incremental.last_value
            self.end_value = incremental.end_value
            self.row_order: TSortOrder = self.incremental.row_order
            self.on_cursor_value_missing = self.incremental.on_cursor_value_missing
            self.range_start = self.incremental.range_start
            self.range_end = self.incremental.range_end
        else:
            self.cursor_column = None
            self.last_value = None
            self.end_value = None
            self.row_order = None
            self.on_cursor_value_missing = None
            self.range_start = None
            self.range_end = None

    def _make_query(self) -> SelectAny:
        table = self.table
        query = table.select()

        # use limit step to limit the query
        if self.limit:
            # limit works with chunks (by default)
            limit = self.limit.limit(self.chunk_size)
            if limit is not None:
                query = query.limit(limit)

        if not self.incremental:
            return query  # type: ignore[no-any-return]
        last_value_func = self.incremental.last_value_func

        # generate where
        if last_value_func is max:  # Query ordered and filtered according to last_value function
            filter_op = operator.ge if self.range_start == "closed" else operator.gt
            filter_op_end = operator.lt if self.range_end == "open" else operator.le
        elif last_value_func is min:
            filter_op = operator.le if self.range_start == "closed" else operator.lt
            filter_op_end = operator.gt if self.range_end == "open" else operator.ge
        else:  # Custom last_value, load everything and let incremental handle filtering
            return query  # type: ignore[no-any-return]

        where_clause = True
        if self.last_value is not None:
            where_clause = filter_op(self.cursor_column, self.last_value)
            if self.end_value is not None:
                where_clause = sa.and_(
                    where_clause, filter_op_end(self.cursor_column, self.end_value)
                )

            if self.on_cursor_value_missing == "include":
                where_clause = sa.or_(where_clause, self.cursor_column.is_(None))
        if self.on_cursor_value_missing == "exclude":
            where_clause = sa.and_(where_clause, self.cursor_column.isnot(None))

        if where_clause is not True:
            query = query.where(where_clause)

        # generate order by from declared row order
        order_by = None
        if (self.row_order == "asc" and last_value_func is max) or (
            self.row_order == "desc" and last_value_func is min
        ):
            order_by = self.cursor_column.asc()
        elif (self.row_order == "asc" and last_value_func is min) or (
            self.row_order == "desc" and last_value_func is max
        ):
            order_by = self.cursor_column.desc()
        if order_by is not None:
            query = query.order_by(order_by)

        return query  # type: ignore[no-any-return]

    def make_query(self) -> SelectClause:
        if self.query_adapter_callback:
            try:
                return self.query_adapter_callback(  # type: ignore[call-arg]
                    self._make_query(), self.table, self.incremental, self.engine
                )
            except TypeError as type_err:
                if not is_typeerror_due_to_wrong_call(type_err, self.query_adapter_callback):
                    raise
                return self.query_adapter_callback(  # type: ignore[call-arg]
                    self._make_query(), self.table
                )

        return self._make_query()

    def compile_query(self, query: SelectClause) -> str:
        """Compile a SQLAlchemy query into a SQL string with literal binds.

        Useful for backends that execute raw SQL strings (ConnectorX, ADBC, etc.).

        Raises:
            NotImplementedError: When the query cannot be compiled to a string.
        """
        try:
            return str(query.compile(self.engine, compile_kwargs={"literal_binds": True}))
        except CompileError as ex:
            raise NotImplementedError(
                f"Query for table `{self.table.name}` could not be compiled to string."
                f" If you are on SQLAlchemy 1.4.x, upgrade to 2.x: `{ex}`"
            ) from ex

    def get_connection_url(self) -> str:
        """Return a plain database connection URL derived from the engine.

        Strips the SQLAlchemy driver portion (e.g. ``+psycopg2``) so the URL
        is suitable for non-SQLAlchemy backends such as ConnectorX or ADBC.
        Override this method when a backend requires a different connection
        string format (e.g. MSSQL ODBC → go-mssqldb conversion).
        """
        return self.engine.url._replace(
            drivername=self.engine.url.get_backend_name()
        ).render_as_string(hide_password=False)

    @abstractmethod
    def load_rows(self, backend_kwargs: Dict[str, Any] = None) -> Iterator[TDataItem]:
        """Load rows from the table and yield them as data items.

        Args:
            backend_kwargs: Backend-specific keyword arguments passed through
                from the resource configuration.

        Yields:
            Data items in the format determined by the loader implementation
            (dicts, Arrow tables, DataFrames, etc.).
        """
        ...


class TableLoader(BaseTableLoader):
    """Default table loader using a SQLAlchemy connection.

    Supports ``"sqlalchemy"``, ``"pyarrow"`` and ``"pandas"`` backends.
    Override ``_load_rows`` to customise query execution (e.g. pagination)
    while reusing ``_convert_result`` for backend format conversion.
    """

    def load_rows(self, backend_kwargs: Dict[str, Any] = None) -> Iterator[TDataItem]:
        # make copy of kwargs
        backend_kwargs = dict(backend_kwargs or {})
        query = self.make_query()
        yield from self._load_rows(query, backend_kwargs)

    def _load_rows(
        self, query: SelectClause, backend_kwargs: Dict[str, Any]
    ) -> Iterator[TDataItem]:
        """Execute *query* over a SQLAlchemy connection and yield converted results.

        Override this method to change the execution strategy (e.g. to add
        pagination or custom retry logic).  Use ``_convert_result`` to convert
        each ``CursorResult`` into the appropriate backend format.
        """
        with self.engine.connect() as conn:
            result = conn.execution_options(yield_per=self.chunk_size).execute(query)
            try:
                yield from self._convert_result(result, backend_kwargs)
            finally:
                result.close()

    def _convert_result(self, result: Any, backend_kwargs: Dict[str, Any]) -> Iterator[TDataItem]:
        """Convert a ``CursorResult`` into data items for the configured backend.

        Partitions *result* into chunks of ``chunk_size`` and converts each
        partition according to ``self.backend``:

        * ``"sqlalchemy"`` -- list of dicts
        * ``"pandas"`` -- pandas ``DataFrame``
        * ``"pyarrow"`` -- Arrow table
        """
        # NOTE: cursor returns not normalized column names! may be quite useful
        # in case of Oracle dialect that normalizes columns
        columns = list(result.keys())
        for partition in result.partitions(size=self.chunk_size):
            if self.backend == "sqlalchemy":
                yield [dict(row._mapping) for row in partition]
            elif self.backend == "pandas":
                from dlt.common.libs.pandas_sql import _wrap_result

                df = _wrap_result(
                    partition,
                    columns,
                    **{"dtype_backend": "pyarrow", **backend_kwargs},
                )
                yield df
            elif self.backend == "pyarrow":
                # add columns present in cursor but missing from schema hints
                for col_name in columns:
                    if col_name not in self.columns:
                        self.columns[col_name] = {"name": col_name}
                # build arrow table using cursor column order for correct
                # positional mapping (self.columns may be in schema order
                # after merge_columns reordering)
                cursor_columns = {c: self.columns[c] for c in columns}
                table = row_tuples_to_arrow(
                    partition,
                    columns=cursor_columns,
                    tz=backend_kwargs.get("tz", "UTC"),
                )
                yield table


class ConnectorXTableLoader(BaseTableLoader):
    """Table loader using ConnectorX for data retrieval. Yields Arrow tables."""

    def load_rows(self, backend_kwargs: Dict[str, Any] = None) -> Iterator[TDataItem]:
        # make copy of kwargs
        backend_kwargs = dict(backend_kwargs or {})
        query = self.make_query()
        yield from self._load_rows_connectorx(query, backend_kwargs)

    def _load_rows_connectorx(
        self, query: SelectClause, backend_kwargs: Dict[str, Any]
    ) -> Iterator[TDataItem]:
        try:
            import connectorx as cx
        except ImportError:
            raise MissingDependencyException("Connector X table backend", ["connectorx"])

        import pyarrow as pa
        from dlt.common.libs.pyarrow import cast_date64_columns_to_timestamp

        # default settings
        backend_kwargs = {
            "protocol": "binary",
            **backend_kwargs,
        }

        is_streaming = False
        if "return_type" in backend_kwargs:
            if backend_kwargs["return_type"] == "arrow_stream":
                is_streaming = True
                backend_kwargs["batch_size"] = backend_kwargs.get("batch_size", self.chunk_size)
        else:
            backend_kwargs["return_type"] = "arrow"

        conn = backend_kwargs.pop("conn", self.get_connection_url())
        query_str = self.compile_query(query)
        logger.info(f"Executing query on ConnectorX: {query_str}")

        if is_streaming:
            record_reader = cx.read_sql(conn, query_str, **backend_kwargs)
            for record_batch in record_reader:
                table = pa.Table.from_batches((record_batch,), schema=record_batch.schema)
                yield cast_date64_columns_to_timestamp(_maybe_fix_0000_timezone(table))
        else:
            df = cx.read_sql(conn, query_str, **backend_kwargs)
            if len(df) > self.chunk_size:
                logger.info(
                    f"The size of the dataset being loaded is more than {self.chunk_size},"
                    " consider using streaming mode (see"
                    " https://dlthub.com/docs/dlt-ecosystem/verified-sources/sql_database"
                    "/configuration#connectorx)"
                )
            yield _maybe_fix_0000_timezone(df)


# populate built-in backends
TABLE_LOADER_REGISTRY.update(
    {
        "sqlalchemy": TableLoader,
        "pyarrow": TableLoader,
        "pandas": TableLoader,
        "connectorx": ConnectorXTableLoader,
    }
)


def register_table_loader_backend(
    backend_name: str,
    loader_class: Type[BaseTableLoader],
) -> None:
    """Register a custom table loader backend.

    After registration the *backend_name* can be used as the ``backend``
    argument of ``sql_table`` / ``sql_database``.

    Args:
        backend_name: Name for the backend (used in ``backend`` parameter).
        loader_class: A subclass of ``BaseTableLoader`` that handles data loading.
    """
    if not issubclass(loader_class, BaseTableLoader):
        raise ValueError(
            f"Invalid table loader: `{loader_class.__name__}`. "
            "Must be a subclass of `BaseTableLoader`."
        )
    TABLE_LOADER_REGISTRY[backend_name] = loader_class
    add_value_to_literal(TableBackend, backend_name)


def get_table_loader_class(backend_name: str) -> Type[BaseTableLoader]:
    """Look up the table loader class for the given backend name."""
    try:
        return TABLE_LOADER_REGISTRY[backend_name]
    except KeyError:
        raise ValueErrorWithKnownValues("backend", backend_name, list(TABLE_LOADER_REGISTRY.keys()))


def table_rows(
    engine: Engine,
    table: Union[Table, str],
    metadata: MetaData,
    chunk_size: int,
    backend: TableBackend,
    incremental: Optional[Incremental[Any]],
    table_adapter_callback: TTableAdapter,
    reflection_level: ReflectionLevel,
    backend_kwargs: Dict[str, Any],
    type_adapter_callback: Optional[TTypeAdapter],
    included_columns: Optional[List[str]],
    excluded_columns: Optional[List[str]],
    query_adapter_callback: Optional[TQueryAdapter],
    resolve_foreign_keys: bool,
    table_loader_class: Optional[Type[BaseTableLoader]] = None,
) -> Iterator[TDataItem]:
    resource = None
    limit = None
    try:
        resource = dlt.current.resource()
        limit = resource.limit
        resource_columns_hints_require_data = callable(resource.columns)
        if resource_columns_hints_require_data and backend == "pyarrow":
            table_name = table.name if isinstance(table, Table) else table
            logger.warning(
                f"Dynamic column hints for '{table_name}' cannot be applied with pyarrow "
                "backend. Use static hints (dict/list) to override reflected types."
            )
    except DltException:
        # in old versions of dlt, resource is not available, so we need to reflect the table again
        pass
    if isinstance(table, str):  # Reflection is deferred
        table = Table(
            table,
            metadata,
            autoload_with=engine,
            extend_existing=True,
            resolve_fks=resolve_foreign_keys,
        )
        table = _execute_table_adapter(
            table, table_adapter_callback, included_columns, excluded_columns
        )
        hints = table_to_resource_hints(
            table,
            reflection_level,
            type_adapter_callback,
            backend == "sqlalchemy",  # skip nested types
            resolve_foreign_keys=resolve_foreign_keys,
        )

        # set the primary_key in the incremental
        # TODO: check for primary key in resource._hints
        if incremental and incremental.primary_key is None:
            primary_key = hints["primary_key"]
            if primary_key is not None:
                incremental.primary_key = primary_key

        # Merge resource hints with reflection hints before yielding
        if resource and hints.get("columns") and not callable(resource.columns):
            hints["columns"] = merge_columns(hints["columns"], resource.columns)

        # yield empty record to set hints and create schema
        # Note: Empty list [] will be written as typed-jsonl (object format), but actual
        # data rows will be written in their native format (e.g., parquet for arrow backend).
        yield dlt.mark.with_hints(
            [],
            dlt.mark.make_hints(
                **hints,
            ),
        )
        # Set columns_hints for TableLoader
        columns_hints = hints["columns"]
    else:
        # table was already reflected -> try to use resource hints
        if not resource or callable(resource.columns):
            hints = table_to_resource_hints(
                table,
                reflection_level,
                type_adapter_callback,
                backend == "sqlalchemy",  # skip nested types
                resolve_foreign_keys=resolve_foreign_keys,
            )
            columns_hints = hints["columns"]

        else:
            # take column hints from resource (which includes user applied hints)
            # Handle callable columns hint (can't resolve without data item)
            columns_hints = resource.columns

    # warn about and remove columns_hints entries that don't match any
    # reflected source column.
    source_columns = {c.name for c in table.columns}
    unmatched = [c for c in columns_hints if c not in source_columns]
    if unmatched:
        logger.warning(
            "Resource columns %s for table '%s' do not match any reflected"
            " source columns and will be ignored. Note: source column names"
            " are not normalized at this stage. Use the original database"
            " column names in apply_hints.",
            unmatched,
            table.name,
        )
        columns_hints = {c: columns_hints[c] for c in columns_hints if c not in unmatched}

    # determine loader class from registry
    if table_loader_class is None:
        table_loader_class = get_table_loader_class(backend)

    loader = table_loader_class(
        engine,
        backend,
        table,
        columns_hints,
        incremental=incremental,
        chunk_size=chunk_size,
        limit=limit,
        query_adapter_callback=query_adapter_callback,
    )
    try:
        yield from loader.load_rows(backend_kwargs)
    finally:
        # dispose the engine if created for this particular table
        # NOTE: database wide engines are not disposed, not externally provided
        if getattr(engine, "may_dispose_after_use", False):
            engine.dispose()


def engine_from_credentials(
    credentials: Union[ConnectionStringCredentials, Engine, str],
    may_dispose_after_use: bool = False,
    **backend_kwargs: Any,
) -> Engine:
    if isinstance(credentials, Engine):
        return credentials
    if isinstance(credentials, ConnectionStringCredentials):
        credentials = credentials.to_native_representation()
    engine = create_engine(credentials, **backend_kwargs)
    setattr(engine, "may_dispose_after_use", may_dispose_after_use)  # noqa
    return engine  # type: ignore[no-any-return]


def unwrap_json_connector_x(field: str) -> TDataItem:
    """Creates a transform function to be added with `add_map` that will unwrap JSON columns
    ingested via connectorx. Such columns are additionally quoted and translate SQL NULL to json "null"
    """
    import pyarrow.compute as pc
    import pyarrow as pa

    def _unwrap(table: TDataItem) -> TDataItem:
        col_index = table.column_names.index(field)
        # remove quotes
        column = table[field]  # pc.replace_substring_regex(table[field], '"(.*)"', "\\1")
        # convert json null to null
        column = pc.replace_with_mask(
            column,
            pc.equal(column, "null").combine_chunks(),
            pa.scalar(None, pa.string()),
        )
        return table.set_column(col_index, table.schema.field(col_index), column)

    return _unwrap


def remove_nullability_adapter(table: Table) -> Table:
    """A table adapter that removes nullability from columns."""
    for col in table.columns:
        # subqueries may not have nullable attr
        if hasattr(col, "nullable"):
            col.nullable = None
    return table


def _execute_table_adapter(
    table: Table,
    adapter: Optional[TTableAdapter],
    included_columns: Optional[List[str]],
    excluded_columns: Optional[List[str]],
) -> Table:
    """Executes default table adapter on `table` and then `adapter` if defined"""
    default_table_adapter(table, included_columns, excluded_columns)
    if adapter:
        # backward compat: old adapters do not return a value
        maybe_query = adapter(table)
        if maybe_query is not None:
            # here we ignore that returned table may be a Select (subquery)
            # otherwise typing gets really complicated
            # TODO: maybe type that later
            table = maybe_query  # type: ignore[assignment]

    return table


def _detect_precision_hints_deprecated(value: Optional[bool]) -> None:
    if value is None:
        return

    msg = (
        "`detect_precision_hints` argument is deprecated and will be removed in a future release. "
    )
    if value:
        msg += "Use `reflection_level='full_with_precision'` which has the same effect instead."

    warnings.warn(
        msg,
        DeprecationWarning,
    )


@configspec
class SqlTableResourceConfiguration(BaseConfiguration):
    credentials: Union[ConnectionStringCredentials, Engine, str] = None
    table: str = None
    schema: Optional[str] = None
    incremental: Optional[Incremental] = None  # type: ignore[type-arg]
    chunk_size: int = 50000
    backend: TableBackend = "sqlalchemy"
    detect_precision_hints: Optional[bool] = None
    defer_table_reflect: Optional[bool] = False
    reflection_level: Optional[ReflectionLevel] = "full"
    included_columns: Optional[List[str]] = None
    excluded_columns: Optional[List[str]] = None
    write_disposition: Optional[TWriteDispositionDict] = None
    primary_key: Optional[TColumnNames] = None
    merge_key: Optional[TColumnNames] = None
    engine_kwargs: Optional[Dict[str, Any]] = None


def _oracle_column_reflect_listener(
    inspector: Any, table: Any, column_info: Dict[str, Any]
) -> None:
    """
    SQLAlchemy event listener for `column_reflect` that enforces Oracle NUMBER type
    to translate to python decimal.Decimal.

    Oracle NUMBER may express floating- or fixed-point numbers, but floats
    don't conform to IEEE754 standard, so we're always using "decimal" type
    to preserve values as accurately as possible. SQLAlchemy2 uses different
    logic for determining asdecimal, we're overriding it here.
    """
    column_type = column_info.get("type")
    if isinstance(column_type, ORACLE_NUMBER):
        column_info["type"] = ORACLE_NUMBER(
            precision=column_type.precision,
            scale=column_type.scale,
            asdecimal=True,
        )


def default_engine_adapter_callback(engine: Engine, metadata: MetaData) -> None:
    """Applies default engine adaptations for known dialects.

    For Oracle dialect, registers an event listener on the provided MetaData that forces
    NUMBER columns to be reflected as Python Decimal to preserve numeric precision.

    Args:
        engine: The SQLAlchemy engine to check dialect for.
        metadata: The MetaData instance to register the listener on.
    """
    if isinstance(engine.dialect, OracleDialect):
        sa.event.listen(metadata, "column_reflect", _oracle_column_reflect_listener)
