import json
import logging
from collections import defaultdict
from typing import Any, Dict, List, Optional

import sqlparse
from dateutil import parser
from sshtunnel import open_tunnel

from legion_query_runner import settings

from ..util import utils
from .utils.requests_session import (
    UnacceptableAddressException,
    requests_or_advocate,
    requests_session,
)

logger = logging.getLogger(__name__)

# Valid types of columns returned in results:
TYPE_INTEGER = "integer"
TYPE_FLOAT = "float"
TYPE_BOOLEAN = "boolean"
TYPE_STRING = "string"
TYPE_DATETIME = "datetime"
TYPE_DATE = "date"

SUPPORTED_COLUMN_TYPES = set([TYPE_INTEGER, TYPE_FLOAT, TYPE_BOOLEAN, TYPE_STRING, TYPE_DATETIME, TYPE_DATE])

class JobTimeoutException(Exception):
    """Exception raised when a job times out."""
    pass

def split_sql_statements(query):
    def strip_trailing_comments(stmt):
        idx = len(stmt.tokens) - 1
        while idx >= 0:
            tok = stmt.tokens[idx]
            if tok.is_whitespace or sqlparse.utils.imt(tok, i=sqlparse.sql.Comment, t=sqlparse.tokens.Comment):
                stmt.tokens[idx] = sqlparse.sql.Token(sqlparse.tokens.Whitespace, " ")
            else:
                break
            idx -= 1
        return stmt

    def strip_trailing_semicolon(stmt):
        idx = len(stmt.tokens) - 1
        while idx >= 0:
            tok = stmt.tokens[idx]
            # we expect that trailing comments already are removed
            if not tok.is_whitespace:
                if sqlparse.utils.imt(tok, t=sqlparse.tokens.Punctuation) and tok.value == ";":
                    stmt.tokens[idx] = sqlparse.sql.Token(sqlparse.tokens.Whitespace, " ")
                break
            idx -= 1
        return stmt

    def is_empty_statement(stmt):
        # copy statement object. `copy.deepcopy` fails to do this, so just re-parse it
        st = sqlparse.engine.FilterStack()
        st.stmtprocess.append(sqlparse.filters.StripCommentsFilter())
        stmt = next(st.run(str(stmt)), None)
        if stmt is None:
            return True

        return str(stmt).strip() == ""

    stack = sqlparse.engine.FilterStack()

    result = [stmt for stmt in stack.run(query)]
    result = [strip_trailing_comments(stmt) for stmt in result]
    result = [strip_trailing_semicolon(stmt) for stmt in result]
    result = [str(stmt).strip() for stmt in result if not is_empty_statement(stmt)]

    if len(result) > 0:
        return result

    return [""]  # if all statements were empty - return a single empty statement

def combine_sql_statements(queries):
    return ";\n".join(queries)

def find_last_keyword_idx(parsed_query):
    for i in reversed(range(len(parsed_query.tokens))):
        if parsed_query.tokens[i].ttype in sqlparse.tokens.Keyword:
            return i
    return -1

class InterruptException(Exception):
    pass

class NotSupported(Exception):
    pass

class BaseQueryRunner:
    deprecated = False
    should_annotate_query = True
    noop_query = None
    limit_query = " LIMIT 1000"
    limit_keywords = ["LIMIT", "OFFSET"]
    limit_after_select = False

    def __init__(self, configuration):
        self.syntax = "sql"
        self.configuration = configuration

    @classmethod
    def name(cls):
        return cls.__name__

    @classmethod
    def type(cls):
        return cls.__name__.lower()

    @classmethod
    def enabled(cls):
        return True

    @property
    def host(self):
        """Returns this query runner's configured host.
        This is used primarily for temporarily swapping endpoints when using SSH tunnels to connect to a data source.

        `BaseQueryRunner`'s naïve implementation supports query runner implementations that store endpoints using `host` and `port`
        configuration values. If your query runner uses a different schema (e.g. a web address), you should override this function.
        """
        if "host" in self.configuration:
            return self.configuration["host"]
        else:
            raise NotImplementedError()

    @host.setter
    def host(self, host):
        """Sets this query runner's configured host.
        This is used primarily for temporarily swapping endpoints when using SSH tunnels to connect to a data source.

        `BaseQueryRunner`'s naïve implementation supports query runner implementations that store endpoints using `host` and `port`
        configuration values. If your query runner uses a different schema (e.g. a web address), you should override this function.
        """
        if "host" in self.configuration:
            self.configuration["host"] = host
        else:
            raise NotImplementedError()

    @property
    def port(self):
        """Returns this query runner's configured port.
        This is used primarily for temporarily swapping endpoints when using SSH tunnels to connect to a data source.

        `BaseQueryRunner`'s naïve implementation supports query runner implementations that store endpoints using `host` and `port`
        configuration values. If your query runner uses a different schema (e.g. a web address), you should override this function.
        """
        if "port" in self.configuration:
            return self.configuration["port"]
        else:
            raise NotImplementedError()

    @port.setter
    def port(self, port):
        """Sets this query runner's configured port.
        This is used primarily for temporarily swapping endpoints when using SSH tunnels to connect to a data source.

        `BaseQueryRunner`'s naïve implementation supports query runner implementations that store endpoints using `host` and `port`
        configuration values. If your query runner uses a different schema (e.g. a web address), you should override this function.
        """
        if "port" in self.configuration:
            self.configuration["port"] = port
        else:
            raise NotImplementedError()

    @classmethod
    def configuration_schema(cls):
        return {}

    def annotate_query(self, query, metadata):
        if not self.should_annotate_query:
            return query

        annotation = ", ".join(["{}: {}".format(k, v) for k, v in metadata.items()])
        annotated_query = "/* {} */ {}".format(annotation, query)
        return annotated_query

    def test_connection(self):
        if self.noop_query is None:
            raise NotImplementedError()
        data, error = self.run_query(self.noop_query, None)

        if error is not None:
            raise Exception(error)

    def run_query(self, query, user):
        raise NotImplementedError()

    def fetch_columns(self, columns):
        column_names = set()
        duplicates_counters = defaultdict(int)
        new_columns = []

        for col in columns:
            column_name = col[0]
            while column_name in column_names:
                duplicates_counters[col[0]] += 1
                column_name = "{}{}".format(col[0], duplicates_counters[col[0]])

            column_names.add(column_name)
            new_columns.append({"name": column_name, "friendly_name": column_name, "type": col[1]})

        return new_columns

    def get_schema(self, get_stats=False):
        raise NotSupported()

    def _handle_run_query_error(self, error):
        if error is None:
            return

        logger.error(error)
        raise Exception(f"Error during query execution. Reason: {error}")

    def _run_query_internal(self, query):
        results, error = self.run_query(query, None)

        if error is not None:
            raise Exception("Failed running query [%s]." % query)
        return results["rows"]

    @classmethod
    def to_dict(cls):
        return {
            "name": cls.name(),
            "type": cls.type(),
            "configuration_schema": cls.configuration_schema(),
            **({"deprecated": True} if cls.deprecated else {}),
        }

    @property
    def supports_auto_limit(self):
        return False

    def apply_auto_limit(self, query_text, should_apply_auto_limit):
        return query_text

    def gen_query_hash(self, query_text, set_auto_limit=False):
        query_text = self.apply_auto_limit(query_text, set_auto_limit)
        return utils.gen_query_hash(query_text)

class BaseSQLQueryRunner(BaseQueryRunner):
    def get_schema(self, get_stats=False):
        schema_dict = {}
        self._get_tables(schema_dict)
        if settings.SCHEMA_RUN_TABLE_SIZE_CALCULATIONS and get_stats:
            self._get_tables_stats(schema_dict)
        return list(schema_dict.values())

    def _get_tables(self, schema_dict):
        return []

    def _get_tables_stats(self, tables_dict):
        for t in tables_dict.keys():
            if isinstance(tables_dict[t], dict):
                res = self._run_query_internal("select count(*) as cnt from %s" % t)
                tables_dict[t]["size"] = res[0]["cnt"]

    @property
    def supports_auto_limit(self):
        return True

    def query_is_select_no_limit(self, query):
        parsed_query = sqlparse.parse(query)[0]
        last_keyword_idx = find_last_keyword_idx(parsed_query)
        # Either invalid query or query that is not select
        if last_keyword_idx == -1 or parsed_query.tokens[0].value.upper() != "SELECT":
            return False

        no_limit = parsed_query.tokens[last_keyword_idx].value.upper() not in self.limit_keywords

        return no_limit

    def add_limit_to_query(self, query):
        parsed_query = sqlparse.parse(query)[0]
        limit_tokens = sqlparse.parse(self.limit_query)[0].tokens
        length = len(parsed_query.tokens)
        if not self.limit_after_select:
            if parsed_query.tokens[length - 1].ttype == sqlparse.tokens.Punctuation:
                parsed_query.tokens[length - 1 : length - 1] = limit_tokens
            else:
                parsed_query.tokens += limit_tokens
        else:
            for i in range(length - 1, -1, -1):
                if parsed_query[i].value.upper() == "SELECT":
                    index = parsed_query.token_index(parsed_query[i + 1])
                    parsed_query = sqlparse.sql.Statement(
                        parsed_query.tokens[:index] + limit_tokens + parsed_query.tokens[index:]
                    )
                    break
        return str(parsed_query)

    def apply_auto_limit(self, query_text, should_apply_auto_limit):
        queries = split_sql_statements(query_text)
        if should_apply_auto_limit:
            # we only check for last one in the list because it is the one that we show result
            last_query = queries[-1]
            if self.query_is_select_no_limit(last_query):
                queries[-1] = self.add_limit_to_query(last_query)
        return combine_sql_statements(queries)

class BaseHTTPQueryRunner(BaseQueryRunner):
    should_annotate_query = False
    response_error = "Endpoint returned unexpected status code"
    requires_authentication = False
    requires_url = True
    url_title = "URL base path"
    username_title = "HTTP Basic Auth Username"
    password_title = "HTTP Basic Auth Password"

    @classmethod
    def configuration_schema(cls):
        schema = {
            "type": "object",
            "properties": {
                "url": {"type": "string", "title": cls.url_title},
                "username": {"type": "string", "title": cls.username_title},
                "password": {"type": "string", "title": cls.password_title},
            },
            "secret": ["password"],
            "order": ["url", "username", "password"],
        }

        if cls.requires_url or cls.requires_authentication:
            schema["required"] = []

        if cls.requires_url:
            schema["required"] += ["url"]

        if cls.requires_authentication:
            schema["required"] += ["username", "password"]
        return schema

    def get_auth(self):
        username = self.configuration.get("username")
        password = self.configuration.get("password")
        if username and password:
            return (username, password)
        if self.requires_authentication:
            raise ValueError("Username and Password required")
        else:
            return None

    def get_response(self, url, auth=None, http_method="get", **kwargs):
        # Get authentication values if not given
        if auth is None:
            auth = self.get_auth()

        # Then call requests to get the response from the given endpoint
        # URL optionally, with the additional requests parameters.
        error = None
        response = None
        try:
            response = requests_session.request(http_method, url, auth=auth, **kwargs)
            # Raise a requests HTTP exception with the appropriate reason
            # for 4xx and 5xx response status codes which is later caught
            # and passed back.
            response.raise_for_status()

            # Any other responses (e.g. 2xx and 3xx):
            if response.status_code != 200:
                error = "{} ({}).".format(self.response_error, response.status_code)

        except requests_or_advocate.HTTPError as exc:
            logger.exception(exc)
            error = "Failed to execute query. "
            f"Return Code: {response.status_code} Reason: {response.text}"
        except UnacceptableAddressException as exc:
            logger.exception(exc)
            error = "Can't query private addresses."
        except requests_or_advocate.RequestException as exc:
            # Catch all other requests exceptions and return the error.
            logger.exception(exc)
            error = str(exc)

        # Return response and error.
        return response, error

query_runners = {}

def register(query_runner_class):
    global query_runners
    if query_runner_class.enabled():
        logger.debug(
            "Registering %s (%s) query runner.",
            query_runner_class.name(),
            query_runner_class.type(),
        )
        query_runners[query_runner_class.type()] = query_runner_class
    else:
        logger.debug(
            "%s query runner enabled but not supported, not registering. Either disable or install missing "
            "dependencies.",
            query_runner_class.name(),
        )

def get_query_runner(query_runner_type, configuration):
    query_runner_class = query_runners.get(query_runner_type, None)
    if query_runner_class is None:
        return None

    return query_runner_class(configuration)

def get_configuration_schema_for_query_runner_type(query_runner_type):
    query_runner_class = query_runners.get(query_runner_type, None)
    if query_runner_class is None:
        return None

    return query_runner_class.configuration_schema()

def import_query_runners(query_runner_imports):
    for runner_import in query_runner_imports:
        __import__(runner_import)

def guess_type(value):
    if isinstance(value, bool):
        return TYPE_BOOLEAN
    elif isinstance(value, int):
        return TYPE_INTEGER
    elif isinstance(value, float):
        return TYPE_FLOAT

    return guess_type_from_string(value)

def guess_type_from_string(string_value):
    if string_value == "" or string_value is None:
        return TYPE_STRING

    try:
        int(string_value)
        return TYPE_INTEGER
    except (ValueError, OverflowError):
        pass

    try:
        float(string_value)
        return TYPE_FLOAT
    except (ValueError, OverflowError):
        pass

    if str(string_value).lower() in ("true", "false"):
        return TYPE_BOOLEAN

    try:
        parser.parse(string_value)
        return TYPE_DATETIME
    except (ValueError, OverflowError):
        pass

    return TYPE_STRING

def with_ssh_tunnel(query_runner, details):
    try:
        remote_host, remote_port = query_runner.host, query_runner.port
    except NotImplementedError:
        raise NotImplementedError("SSH tunneling is not implemented for this query runner yet.")

    bastion_address = (details["ssh_host"], details.get("ssh_port", 22))
    remote_address = (remote_host, remote_port)
    auth = {
        "ssh_username": details["ssh_username"],
        **settings.dynamic_settings.ssh_tunnel_auth(),
    }

    try:
        server = open_tunnel(bastion_address, remote_bind_address=remote_address, **auth)
        server = server.__enter__()
        if not hasattr(server, 'local_bind_address') or not server.local_bind_address:
            raise ValueError("SSH tunnel server did not provide local bind address")
        query_runner.host, query_runner.port = server.local_bind_address
        return server
    except Exception as error:
        raise type(error)("SSH tunnel: {}".format(str(error)))

class QueryRunner:
    """A unified adapter for different query runners that provides a consistent interface."""

    def __init__(self, db_type: str, configuration: Dict[str, Any]):
        """
        Initialize the query runner adapter.

        Args:
            db_type: Type of database ('mysql', 'pg', 'redshift', etc.)
            configuration: Database connection configuration
        """
        self.runner = get_query_runner(db_type, configuration)
        if not self.runner:
            raise ValueError(f"Unsupported database type: {db_type}")

    def test_connection(self) -> None:
        """Test the database connection."""
        self.runner.test_connection()

    def run_query(self, query: str, user: Optional[Any] = None) -> Dict[str, Any]:
        """
        Run a query and return the results.

        Args:
            query: SQL query to execute
            user: Optional user context

        Returns:
            Dict containing 'columns' and 'rows' if successful

        Raises:
            Exception: If the query fails or SSH tunnel fails
        """
        data, error = self.runner.run_query(query, user)
        if error:
            raise Exception(error)
        if not data:
            return {'columns': [], 'rows': []}
        if isinstance(data, str):
            return json.loads(data)
        return data

    def get_schema(self, get_stats: bool = False) -> List[Dict[str, Any]]:
        """
        Get the database schema.

        Args:
            get_stats: Whether to include table statistics

        Returns:
            List of table definitions with their columns
        """
        return self.runner.get_schema(get_stats)

    def get_table_columns(self, table_name: str) -> List[str]:
        """
        Get column names for a specific table.

        Args:
            table_name: Name of the table

        Returns:
            List of column names
        """
        schema = self.get_schema()
        for table in schema:
            if table['name'] == table_name:
                return [col['name'] if isinstance(col, dict) else col for col in table['columns']]
        return []

    def get_table_types(self, table_name: str) -> Dict[str, Optional[str]]:
        """
        Get column types for a specific table.

        Args:
            table_name: Name of the table

        Returns:
            Dict mapping column names to their types (None if column is a string)
        """
        schema = self.get_schema()
        for table in schema:
            if table['name'] == table_name:
                return {
                    col['name'] if isinstance(col, dict) else col:
                    col.get('type', TYPE_STRING) if isinstance(col, dict) else None
                    for col in table['columns']
                }
        return {}
