import datetime
import logging
import os

from legion_query_runner.settings.helpers import cast_int_or_default

from .base import (
    TYPE_BOOLEAN,
    TYPE_DATE,
    TYPE_DATETIME,
    TYPE_FLOAT,
    TYPE_INTEGER,
    TYPE_STRING,
    BaseSQLQueryRunner,
    register,
    split_sql_statements,
)

try:
    from databricks import sql

    enabled = True
except ImportError:
    enabled = False

TYPES_MAP = {
    str: TYPE_STRING,
    bool: TYPE_BOOLEAN,
    datetime.date: TYPE_DATE,
    datetime.datetime: TYPE_DATETIME,
    int: TYPE_INTEGER,
    float: TYPE_FLOAT,
}

ROW_LIMIT = cast_int_or_default(os.environ.get("DATABRICKS_ROW_LIMIT"), 20000)

logger = logging.getLogger(__name__)


class Databricks(BaseSQLQueryRunner):
    noop_query = "SELECT 1"
    should_annotate_query = False

    @classmethod
    def type(cls):
        return "databricks"

    @classmethod
    def enabled(cls):
        return enabled

    @classmethod
    def configuration_schema(cls):
        return {
            "type": "object",
            "properties": {
                "host": {"type": "string", "title": "Server Hostname"},
                "http_path": {"type": "string", "title": "HTTP Path"},
                "http_password": {"type": "string", "title": "Access Token"},
                # Legacy support - map http_password to access_token
                "access_token": {"type": "string", "title": "Access Token"},
            },
            "order": ["host", "http_path", "http_password"],
            "secret": ["http_password", "access_token"],
            "required": ["host", "http_path"],
        }

    def _get_connection(self):
        """Create a connection to Databricks using the native SQL connector"""
        host = self.configuration["host"]
        http_path = self.configuration["http_path"]
        # Support both http_password (legacy) and access_token
        access_token = self.configuration.get("access_token") or self.configuration.get("http_password")

        if not access_token:
            raise ValueError("access_token or http_password is required")

        return sql.connect(
            server_hostname=host,
            http_path=http_path,
            access_token=access_token,
        )

    def run_query(self, query, user):
        connection = None
        cursor = None
        try:
            connection = self._get_connection()
            cursor = connection.cursor()

            statements = split_sql_statements(query)
            for stmt in statements:
                cursor.execute(stmt)

            if cursor.description is not None:
                # Fetch all rows (up to ROW_LIMIT)
                result_set = cursor.fetchmany(ROW_LIMIT)

                # Get column information from cursor.description
                # cursor.description is a list of tuples: (name, type_code, display_size, internal_size, precision, scale, null_ok)
                columns = self.fetch_columns([
                    (desc[0], None)  # Column name, type (we'll infer from data if needed)
                    for desc in cursor.description
                ])

                rows = [dict(zip((column["name"] for column in columns), row)) for row in result_set]

                data = {"columns": columns, "rows": rows}

                # Check if there are more rows
                if len(result_set) >= ROW_LIMIT:
                    # Try to fetch one more to see if there are more results
                    try:
                        next_row = cursor.fetchone()
                        if next_row is not None:
                            logger.warning("Truncated result set.")
                            data["truncated"] = True
                    except Exception:
                        # If fetchone fails, assume we got all results
                        pass

                error = None
            else:
                error = None
                data = {
                    "columns": [{"name": "result", "type": TYPE_STRING}],
                    "rows": [{"result": "No data was returned."}],
                }

        except Exception as e:
            error = str(e)
            data = None
            logger.error(f"Databricks query error: {error}", exc_info=True)
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

        return data, error

    def test_connection(self):
        """Test the database connection"""
        connection = None
        try:
            connection = self._get_connection()
            cursor = connection.cursor()
            cursor.execute(self.noop_query)
            cursor.fetchone()
            cursor.close()
        finally:
            if connection:
                connection.close()

    def get_schema(self, get_stats=False):
        """Get the database schema by iterating through all databases and their tables."""
        schema = {}
        databases = self.get_databases()
        for database in databases:
            tables = self.get_database_tables_with_columns(database)
            for table in tables:
                schema[table['name']] = table
        return list(schema.values())

    def get_databases(self):
        query = "SHOW DATABASES"
        results, error = self.run_query(query, None)

        if error is not None:
            self._handle_run_query_error(error)

        first_column_name = results["columns"][0]["name"]
        return [row[first_column_name] for row in results["rows"]]

    def get_database_tables(self, database_name):
        # Use SHOW TABLES instead of ODBC cursor.tables()
        query = f"SHOW TABLES IN {database_name}"
        results, error = self.run_query(query, None)

        if error is not None:
            self._handle_run_query_error(error)

        schema = []
        # The result format depends on Databricks version
        # Typically: [{"database": "...", "tableName": "...", ...}]
        for row in results["rows"]:
            table_name = row.get("tableName") or row.get("table") or list(row.values())[0]
            schema.append({"name": f"{database_name}.{table_name}", "columns": []})

        return schema

    def get_database_tables_with_columns(self, database_name):
        schema = {}

        # Get all tables
        tables_result, error = self.run_query(f"SHOW TABLES IN {database_name}", None)
        if error is not None:
            self._handle_run_query_error(error)

        for row in tables_result["rows"]:
            table_name = row.get("tableName") or row.get("table") or list(row.values())[0]
            full_table_name = f"{database_name}.{table_name}"
            schema[full_table_name] = {"name": full_table_name, "columns": []}

        # Get columns for each table
        for table_name in schema.keys():
            db_name, tbl_name = table_name.split(".", 1)
            try:
                columns_result, error = self.run_query(
                    f"DESCRIBE TABLE {db_name}.{tbl_name}", None
                )
                if error is None:
                    for col_row in columns_result["rows"]:
                        col_name = col_row.get("col_name") or col_row.get("name") or list(col_row.values())[0]
                        col_type = col_row.get("data_type") or col_row.get("type") or list(col_row.values())[1] if len(col_row) > 1 else "string"
                        schema[table_name]["columns"].append({"name": col_name, "type": col_type})
            except Exception as e:
                logger.warning(f"Could not get columns for {table_name}: {e}")

        return list(schema.values())

    def get_table_columns(self, database_name, table_name):
        results, error = self.run_query(
            f"DESCRIBE TABLE {database_name}.{table_name}", None
        )
        if error is not None:
            self._handle_run_query_error(error)

        columns = []
        for row in results["rows"]:
            col_name = row.get("col_name") or row.get("name") or list(row.values())[0]
            col_type = row.get("data_type") or row.get("type") or list(row.values())[1] if len(row) > 1 else "string"
            columns.append({"name": col_name, "type": col_type})

        return columns


register(Databricks)
