# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

from contextlib import nullcontext
from typing import TYPE_CHECKING, Any, cast

from marimo import _loggers
from marimo._config.config import SqlOutputType
from marimo._data.models import DataType
from marimo._dependencies.dependencies import DependencyManager
from marimo._runtime._wasm._duckdb import (
    try_run_duckdb_sql_with_wasm_patch,
)
from marimo._runtime.context.types import (
    ContextNotInitializedError,
    get_context,
    runtime_context_installed,
)

if TYPE_CHECKING:
    from collections.abc import Callable

    import duckdb
    import pandas as pd
    import polars as pl
    from polars._typing import ConnectionOrCursor

LOGGER = _loggers.marimo_logger()

CHEAP_DISCOVERY_DATABASES = {
    "duckdb",
    "sqlite",
    "mysql",
    "mariadb",
    "postgresql",
}
# DuckDB SQL can return None for DDL, so keep "patch did not apply" distinct.
_NO_WASM_DUCKDB_RESULT = object()


def is_cheap_dialect(dialect: str) -> bool:
    """Whether schema/table discovery is cheap for the given SQL dialect.

    Used as the default heuristic for resolving `"auto"` discovery flags: local
    or lightweight dialects are cheap to introspect, whereas remote analytical
    backends (e.g. Snowflake, BigQuery) can be slow and should not be scanned
    eagerly.
    """
    return dialect.lower() in CHEAP_DISCOVERY_DATABASES


def get_configured_sql_output_format() -> SqlOutputType:
    """Read the configured SQL output format from the runtime context.

    Returns "auto" when no runtime context is available (e.g. when `mo.sql(...)`
    is called outside of a marimo app).
    """
    if not runtime_context_installed():
        return "auto"
    try:
        return get_context().app_config.sql_output
    except ContextNotInitializedError:
        return "auto"


def _try_wasm_duckdb(
    method_name: str,
    query: str,
    connection: Any,
    glbls: dict[str, Any],
    *trailing_args: Any,
) -> object:
    import duckdb

    if connection is duckdb:
        original = getattr(duckdb, method_name)
        args: tuple[Any, ...] = (query, *trailing_args)
        query_arg_index = 0
    else:
        original = getattr(type(connection), method_name)
        args = (connection, query, *trailing_args)
        query_arg_index = 1

    result = try_run_duckdb_sql_with_wasm_patch(
        original,
        args,
        {},
        query_arg_index=query_arg_index,
        query_kwarg_names=("query",),
        eval_globals=glbls,
        eval_locals=glbls,
    )
    return _NO_WASM_DUCKDB_RESULT if result is None else result.value


def wrapped_sql(
    query: str,
    connection: duckdb.DuckDBPyConnection | None,
) -> duckdb.DuckDBPyRelation:
    DependencyManager.duckdb.require("to execute sql")
    import duckdb

    # In Python globals() are scoped to modules; since this function
    # is in a different module than user code, globals() doesn't return
    # the kernel globals, it just returns this module's global namespace.
    #
    # However, duckdb needs access to the kernel's globals. For this reason,
    # we manually exec duckdb and provide it with the kernel's globals.
    if connection is None:
        connection = cast(duckdb.DuckDBPyConnection, duckdb)

    try:
        ctx = get_context()
    except ContextNotInitializedError:
        result = _try_wasm_duckdb("sql", query, connection, globals())
        if result is _NO_WASM_DUCKDB_RESULT:
            # No WASM rewrite was needed; use DuckDB's normal SQL path.
            relation = connection.sql(query=query)
        else:
            relation = cast(duckdb.DuckDBPyRelation, result)
    else:
        install_connection = (
            ctx.execution_context.with_connection
            if ctx.execution_context is not None
            else nullcontext
        )
        with install_connection(connection):
            result = _try_wasm_duckdb(
                "sql",
                query,
                connection,
                ctx.globals,
            )
            if result is _NO_WASM_DUCKDB_RESULT:
                # Run in kernel globals so DuckDB replacement scans see user data.
                relation = eval(
                    "connection.sql(query=query)",
                    ctx.globals,
                    {"query": query, "connection": connection},
                )
            else:
                relation = cast(duckdb.DuckDBPyRelation, result)

    return relation


def execute_duckdb_sql(
    query: str,
    params: list[Any],
    connection: duckdb.DuckDBPyConnection | None = None,
) -> duckdb.DuckDBPyConnection:
    """Execute a parameterized DuckDB query with kernel globals context.

    Like wrapped_sql, but uses connection.execute() to support
    parameterized queries ($1, $2, ...) for safe value interpolation.
    """
    DependencyManager.duckdb.require("to execute sql")
    import duckdb

    if connection is None:
        connection = cast(duckdb.DuckDBPyConnection, duckdb)

    try:
        ctx = get_context()
    except ContextNotInitializedError:
        result = _try_wasm_duckdb(
            "execute", query, connection, globals(), params
        )
        if result is _NO_WASM_DUCKDB_RESULT:
            # No WASM rewrite was needed; preserve DuckDB's parameterized path.
            return connection.execute(query, params)
        return cast(duckdb.DuckDBPyConnection, result)
    else:
        install_connection = (
            ctx.execution_context.with_connection
            if ctx.execution_context is not None
            else nullcontext
        )
        with install_connection(connection):
            result = _try_wasm_duckdb(
                "execute",
                query,
                connection,
                ctx.globals,
                params,
            )
            if result is _NO_WASM_DUCKDB_RESULT:
                # Run in kernel globals so parameterized SQL can scan user data.
                value = eval(
                    "connection.execute(query, params)",
                    ctx.globals,
                    {
                        "query": query,
                        "params": params,
                        "connection": connection,
                    },
                )
            else:
                value = result
            return cast(duckdb.DuckDBPyConnection, value)


def try_convert_to_polars(
    *,
    query: str,
    connection: ConnectionOrCursor,
    lazy: bool,
) -> tuple[pl.DataFrame | pl.LazyFrame | None, Exception | None]:
    """Try to convert the query to a polars dataframe.

    Returns:
        - The polars dataframe, or None if the conversion failed.
        - Error message, or None if the conversion succeeded.
    """
    import polars as pl

    try:
        df = pl.read_database(query=query, connection=connection)
        return df.lazy() if lazy else df, None
    except (
        pl.exceptions.PanicException,
        pl.exceptions.ComputeError,
    ) as e:
        return None, e


def convert_to_output(
    *,
    sql_output_format: SqlOutputType,
    to_polars: Callable[[], pl.DataFrame | pl.Series],
    to_pandas: Callable[[], pd.DataFrame],
    to_native: Callable[[], Any] | None = None,
    to_lazy_polars: Callable[[], pl.LazyFrame] | None = None,
) -> Any:
    """Convert a result to the specified output format.

    Args:
        sql_output_format (SqlOutputType): The output format to convert to.
        to_polars (Callable[[], Any]): A function to convert the result to polars.
        to_pandas (Callable[[], Any]): A function to convert the result to pandas.
        to_native (Callable[[], Any]): A function to convert the result to native.
        to_lazy_polars (Optional[Callable[[], Any]]): A function to convert the result to lazy polars.
            If not present, will use to_polars().lazy()

    Returns:
        Any: The converted result.
    """
    if sql_output_format == "native":
        if to_native is None:
            raise ValueError("to_native is required for native output format")
        return to_native()

    if sql_output_format in ("polars", "lazy-polars"):
        if not DependencyManager.polars.has():
            raise_df_import_error("polars[pyarrow]")

    if sql_output_format == "polars":
        return to_polars()

    if sql_output_format == "lazy-polars":
        import polars as pl

        if to_lazy_polars is not None:
            return to_lazy_polars()

        # Default handling, we convert to polars and then to lazy polars
        result = to_polars()
        if isinstance(result, pl.Series):
            return result.to_frame().lazy()
        return result.lazy()

    if sql_output_format == "pandas":
        if not DependencyManager.pandas.has():
            raise_df_import_error("pandas")
        return to_pandas()

    # Auto
    if DependencyManager.polars.has():
        import polars as pl

        try:
            return to_polars()
        except (
            pl.exceptions.PanicException,
            pl.exceptions.ComputeError,
        ):
            LOGGER.info("Failed to convert to polars, falling back to pandas")
            DependencyManager.pandas.require("to convert this data")

    if DependencyManager.pandas.has():
        try:
            return to_pandas()
        except Exception as e:
            LOGGER.warning("Failed to convert dataframe", exc_info=e)
            return None

    raise_df_import_error("polars[pyarrow]")


def raise_df_import_error(pkg: str) -> None:
    raise ModuleNotFoundError(
        "pandas or polars is required to execute sql. "
        + "You can install them with 'pip install pandas polars'",
        name=pkg,
    )


def sql_type_to_data_type(type_str: str) -> DataType:
    """Convert SQL type string to DataType"""
    type_str = type_str.lower()
    if any(x in type_str for x in ("int", "serial")):
        return "integer"
    elif any(x in type_str for x in ("float", "double", "decimal", "numeric")):
        return "number"
    elif any(x in type_str for x in ("timestamp", "datetime")):
        return "datetime"
    elif "date" in type_str:
        return "date"
    elif "time" in type_str:
        return "time"
    elif "bool" in type_str:
        return "boolean"
    elif any(x in type_str for x in ("char", "text")):
        return "string"
    else:
        return "string"


def is_explain_query(query: str) -> bool:
    """Check if a SQL query is an EXPLAIN query."""
    import re

    return bool(re.match(r"\s*explain\s", query, re.IGNORECASE))


def wrap_query_with_explain(query: str) -> str:
    """
    Wrap a SQL query with an EXPLAIN query if it is not already.

    If the query is just comments, return it. Executing this would return nothing.
    """
    if is_explain_query(query):
        return query

    return f"EXPLAIN {query}"


def is_query_empty(query: str) -> bool:
    """Check if a SQL query is empty or just comments"""
    stripped = query.strip()
    if not stripped:
        return True

    # If the query starts with -- or /*, it's likely just comments
    if stripped.startswith(("--", "/*")):
        import re

        # Remove /* */ comments
        no_block_comments = re.sub(r"/\*.*?\*/", "", query, flags=re.DOTALL)

        # Remove -- comments (just split on \n and check each line)
        lines = no_block_comments.split("\n")
        for line in lines:
            # Find first non-whitespace character
            for char in line:
                if char.isspace():
                    continue
                elif char == "-":
                    if line.strip().startswith("--"):
                        break  # This line is a comment, continue to next line
                    else:
                        return False  # Found non-comment content
                else:
                    return False  # Found non-comment content
        return True

    # If it doesn't start with comment markers, it's not empty
    return False


def extract_explain_content(df: Any) -> str:
    """Extract all content from a DataFrame for EXPLAIN queries.

    Args:
        df: DataFrame (pandas or polars). If not pandas / polars / duckdb relation, return repr(df).

    Returns:
        String containing content of dataframe
    """
    try:
        if DependencyManager.polars.imported():
            import polars as pl

            if isinstance(df, pl.LazyFrame):
                df = df.collect()
            if isinstance(df, pl.DataFrame):
                # Display full strings without truncation
                with pl.Config(fmt_str_lengths=10000):
                    return str(df)

        if DependencyManager.pandas.imported():
            import pandas as pd

            if isinstance(df, pd.DataFrame):
                # Preserve newlines in the data
                all_values = df.values.flatten().tolist()
                return "\n".join(str(val) for val in all_values)

        if DependencyManager.duckdb.imported():
            import duckdb

            if isinstance(df, duckdb.DuckDBPyRelation):
                rows = df.fetchall()
                return "\n".join(str(val) for row in rows for val in row)

        # Fallback to repr for other types
        return repr(df)

    except Exception as e:
        LOGGER.debug("Failed to extract explain content: %s", e)
        return repr(df)


def strip_explain_from_error_message(error_message: str) -> str:
    """Strip EXPLAIN from an error message. Also adjusts the caret position

    Example:
    ```
    LINE 1: EXPLAIN SELECT * FROM t
                                  ^
    ```
    becomes
    ```
    LINE 1: SELECT * FROM t
                          ^
    ```
    """
    # Find the first occurrence of "EXPLAIN " and replace it
    explain_pos = error_message.find("EXPLAIN ")
    if explain_pos == -1:
        return error_message

    explain_length = len("EXPLAIN ")

    # Replace the first "EXPLAIN " with empty string
    result = (
        error_message[:explain_pos]
        + error_message[explain_pos + explain_length :]
    )

    # Find the next newline and strip the same amount from the next line
    next_newline = result.find("\n", explain_pos)
    if next_newline != -1 and next_newline + explain_length < len(result):
        # Remove the same length from the beginning of the next line
        # This is the caret position
        result = (
            result[: next_newline + 1]
            + result[next_newline + 1 + explain_length :]
        )

    return result
