#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import json
import logging
from json import JSONDecodeError
from typing import Any, overload
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlsplit

import attrs

from airflow.sdk.exceptions import AirflowException, AirflowNotFoundException, AirflowRuntimeError, ErrorType
from airflow.sdk.providers_manager_runtime import ProvidersManagerTaskRuntime

log = logging.getLogger(__name__)


def _parse_netloc_to_hostname(uri_parts):
    """
    Parse a URI string to get the correct Hostname.

    ``urlparse(...).hostname`` or ``urlsplit(...).hostname`` returns value into the lowercase in most cases,
    there are some exclusion exists for specific cases such as https://bugs.python.org/issue32323
    In case if expected to get a path as part of hostname path,
    then default behavior ``urlparse``/``urlsplit`` is unexpected.
    """
    hostname = unquote(uri_parts.hostname or "")
    if "/" in hostname:
        hostname = uri_parts.netloc
        if "@" in hostname:
            hostname = hostname.rsplit("@", 1)[1]
        if ":" in hostname:
            hostname = hostname.split(":", 1)[0]
        hostname = unquote(hostname)
    return hostname


def _prune_dict(val: Any, mode="strict"):
    """
    Given dict ``val``, returns new dict based on ``val`` with all empty elements removed.

    What constitutes "empty" is controlled by the ``mode`` parameter.  If mode is 'strict'
    then only ``None`` elements will be removed.  If mode is ``truthy``, then element ``x``
    will be removed if ``bool(x) is False``.
    """
    if mode == "strict":
        is_empty = lambda x: x is None
    elif mode == "truthy":
        is_empty = lambda x: not bool(x)
    else:
        raise ValueError("allowable values for `mode` include 'truthy' and 'strict'")

    if isinstance(val, dict):
        new_dict = {}
        for k, v in val.items():
            if is_empty(v):
                continue
            if isinstance(v, (list, dict)):
                new_val = _prune_dict(v, mode=mode)
                if not is_empty(new_val):
                    new_dict[k] = new_val
            else:
                new_dict[k] = v
        return new_dict
    if isinstance(val, list):
        new_list = []
        for v in val:
            if is_empty(v):
                continue
            if isinstance(v, (list, dict)):
                new_val = _prune_dict(v, mode=mode)
                if not is_empty(new_val):
                    new_list.append(new_val)
            else:
                new_list.append(v)
        return new_list
    return val


@attrs.define(slots=False)
class Connection:
    """
    A connection to an external data source.

    :param conn_id: The connection ID.
    :param conn_type: The connection type.
    :param description: The connection description.
    :param host: The host.
    :param login: The login.
    :param password: The password.
    :param schema: The schema.
    :param port: The port number.
    :param extra: Extra metadata. Non-standard data such as private/SSH keys can be saved here. JSON
        encoded object.
    :param uri: URI address describing connection parameters.
    """

    conn_id: str
    conn_type: str | None = None
    description: str | None = None
    host: str | None = None
    schema: str | None = None
    login: str | None = None
    password: str | None = None
    port: int | None = None
    extra: str | None = None

    EXTRA_KEY = "__extra__"

    @overload
    def __init__(self, *, conn_id: str, uri: str) -> None: ...

    @overload
    def __init__(
        self,
        *,
        conn_id: str,
        conn_type: str | None = None,
        description: str | None = None,
        host: str | None = None,
        schema: str | None = None,
        login: str | None = None,
        password: str | None = None,
        port: int | None = None,
        extra: str | None = None,
    ) -> None: ...

    def __init__(self, *, conn_id: str, uri: str | None = None, **kwargs) -> None:
        if uri is not None and kwargs:
            raise AirflowException(
                "You must create an object using the URI or individual values "
                "(conn_type, host, login, password, schema, port or extra). "
                "You can't mix these two ways to create this object."
            )
        if uri is None:
            self.__attrs_init__(conn_id=conn_id, **kwargs)  # type: ignore[attr-defined]
        else:
            self.__dict__.update(attrs.asdict(self.from_uri(uri, conn_id=conn_id), recurse=False))

    def get_uri(self) -> str:
        """Generate and return connection in URI format."""
        from urllib.parse import parse_qsl

        if self.conn_type:
            uri = f"{self.conn_type.lower().replace('_', '-')}://"
        else:
            uri = "//"
        host_to_use: str | None
        if self.host and "://" in self.host:
            protocol, host = self.host.split("://", 1)
            # If the protocol in host matches the connection type, don't add it again
            if protocol == self.conn_type:
                host_to_use = self.host
                protocol_to_add = None
            else:
                # Different protocol, add it to the URI
                host_to_use = host
                protocol_to_add = protocol
        else:
            host_to_use = self.host
            protocol_to_add = None

        if protocol_to_add:
            uri += f"{protocol_to_add}://"

        authority_block = ""
        if self.login is not None:
            authority_block += quote(self.login, safe="")
        if self.password is not None:
            authority_block += ":" + quote(self.password, safe="")
        if authority_block > "":
            authority_block += "@"
            uri += authority_block

        host_block = ""
        if host_to_use:
            host_block += quote(host_to_use, safe="")
        if self.port:
            if host_block == "" and authority_block == "":
                host_block += f"@:{self.port}"
            else:
                host_block += f":{self.port}"
        if self.schema:
            host_block += f"/{quote(self.schema, safe='')}"
        uri += host_block

        if self.extra:
            try:
                extra_dejson = self.extra_dejson
                query: str | None = urlencode(extra_dejson)
            except TypeError:
                query = None
            if query and extra_dejson == dict(parse_qsl(query, keep_blank_values=True)):
                uri += ("?" if self.schema else "/?") + query
            else:
                uri += ("?" if self.schema else "/?") + urlencode({self.EXTRA_KEY: self.extra})

        return uri

    def get_hook(self, *, hook_params=None):
        """Return hook based on conn_type."""
        from airflow.sdk._shared.module_loading import import_string

        hook = ProvidersManagerTaskRuntime().hooks.get(self.conn_type, None)

        if hook is None:
            raise AirflowException(f'Unknown hook type "{self.conn_type}"')
        try:
            hook_class = import_string(hook.hook_class_name)
        except ImportError:
            log.error(
                "Could not import %s when discovering %s %s",
                hook.hook_class_name,
                hook.hook_name,
                hook.package_name,
            )
            raise
        if hook_params is None:
            hook_params = {}
        return hook_class(**{hook.connection_id_attribute_name: self.conn_id}, **hook_params)

    def test_connection(self) -> tuple[bool, str]:
        """
        Call ``get_hook`` and execute ``test_connection`` on the resulting hook.

        Pre-warms ``_preset_connections`` with ``self`` so the hook can resolve
        this connection by ``conn_id`` from inside ``hook.test_connection`` even
        when no metadata-DB or secrets backend has it yet (used by the async
        connection-test workflow where the worker holds the connection in memory).
        """
        from airflow.sdk.execution_time.context import _preset_connections

        outer = _preset_connections.get() or {}
        reset_token = _preset_connections.set({**outer, self.conn_id: self})
        try:
            try:
                hook = self.get_hook()
            except AirflowException as e:
                return False, str(e)
            if not getattr(hook, "test_connection", None):
                return (
                    False,
                    f"Hook {type(hook).__name__} doesn't implement or inherit test_connection method",
                )
            try:
                return hook.test_connection()
            except Exception as e:
                return False, str(e)
        finally:
            _preset_connections.reset(reset_token)

    @classmethod
    def _handle_connection_error(cls, e: AirflowRuntimeError, conn_id: str) -> None:
        """Handle connection retrieval errors."""
        if e.error.error == ErrorType.CONNECTION_NOT_FOUND:
            raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't defined") from None
        raise

    @classmethod
    def get(cls, conn_id: str) -> Any:
        from airflow.sdk.execution_time.context import _get_connection

        try:
            return _get_connection(conn_id)
        except AirflowRuntimeError as e:
            cls._handle_connection_error(e, conn_id)

    @classmethod
    async def async_get(cls, conn_id: str) -> Any:
        from airflow.sdk.execution_time.context import _async_get_connection

        try:
            return await _async_get_connection(conn_id)
        except AirflowRuntimeError as e:
            cls._handle_connection_error(e, conn_id)

    @property
    def extra_dejson(self) -> dict:
        """Returns the extra property by deserializing json."""
        from airflow.sdk.log import mask_secret

        extra = {}
        if self.extra:
            try:
                extra = json.loads(self.extra)
            except JSONDecodeError:
                log.exception("Failed to deserialize extra property `extra`, returning empty dictionary")
            else:
                mask_secret(extra)

        return extra

    def get_extra_dejson(self) -> dict:
        """Deserialize extra property to JSON."""
        import warnings

        warnings.warn(
            "`get_extra_dejson` is deprecated and will be removed in a future release. ",
            DeprecationWarning,
            stacklevel=2,
        )
        return self.extra_dejson

    def to_dict(self, *, prune_empty: bool = False, validate: bool = True) -> dict[str, Any]:
        """
        Convert Connection to json-serializable dictionary.

        :param prune_empty: Whether or not remove empty values.
        :param validate: Validate dictionary is JSON-serializable

        :meta private:
        """
        conn: dict[str, Any] = {
            "conn_id": self.conn_id,
            "conn_type": self.conn_type,
            "description": self.description,
            "host": self.host,
            "login": self.login,
            "password": self.password,
            "schema": self.schema,
            "port": self.port,
        }
        if prune_empty:
            conn = _prune_dict(val=conn, mode="strict")
        if (extra := self.extra_dejson) or not prune_empty:
            conn["extra"] = extra

        if validate:
            json.dumps(conn)
        return conn

    @classmethod
    def from_json(cls, value, conn_id=None) -> Connection:
        kwargs = json.loads(value)
        extra = kwargs.pop("extra", None)
        if extra:
            kwargs["extra"] = extra if isinstance(extra, str) else json.dumps(extra)
        conn_type = kwargs.pop("conn_type", None)
        if conn_type:
            kwargs["conn_type"] = cls._normalize_conn_type(conn_type)
        port = kwargs.pop("port", None)
        if port:
            try:
                kwargs["port"] = int(port)
            except ValueError:
                raise ValueError(f"Expected integer value for `port`, but got {port!r} instead.")
        return cls(conn_id=conn_id, **kwargs)

    def as_json(self) -> str:
        """Convert Connection to JSON-string object."""
        conn_repr = self.to_dict(prune_empty=True, validate=False)
        conn_repr.pop("conn_id", None)
        return json.dumps(conn_repr)

    @classmethod
    def from_uri(cls, uri: str, conn_id: str) -> Connection:
        """
        Create a Connection from a URI string.

        :param uri: URI string to parse
        :param conn_id: Connection ID to assign to the connection
        :return: Connection object
        """
        schemes_count_in_uri = uri.count("://")
        if schemes_count_in_uri > 2:
            raise AirflowException(f"Invalid connection string: {uri}.")
        host_with_protocol = schemes_count_in_uri == 2
        uri_parts = urlsplit(uri)
        conn_type = uri_parts.scheme
        normalized_conn_type = cls._normalize_conn_type(conn_type)
        rest_of_the_url = uri.replace(f"{conn_type}://", ("" if host_with_protocol else "//"))
        if host_with_protocol:
            uri_splits = rest_of_the_url.split("://", 1)
            if "@" in uri_splits[0] or ":" in uri_splits[0]:
                raise AirflowException(f"Invalid connection string: {uri}.")
        uri_parts = urlsplit(rest_of_the_url)
        protocol = uri_parts.scheme if host_with_protocol else None
        host = _parse_netloc_to_hostname(uri_parts)
        parsed_host = cls._create_host(protocol, host)
        quoted_schema = uri_parts.path[1:]
        schema = unquote(quoted_schema) if quoted_schema else quoted_schema
        login = unquote(uri_parts.username) if uri_parts.username else uri_parts.username
        password = unquote(uri_parts.password) if uri_parts.password else uri_parts.password
        port = uri_parts.port
        extra = None
        if uri_parts.query:
            query = dict(parse_qsl(uri_parts.query, keep_blank_values=True))
            if cls.EXTRA_KEY in query:
                extra = query[cls.EXTRA_KEY]
            else:
                extra = json.dumps(query)

        return cls(
            conn_id=conn_id,
            conn_type=normalized_conn_type,
            host=parsed_host,
            schema=schema,
            login=login,
            password=password,
            port=port,
            extra=extra,
        )

    @staticmethod
    def _create_host(protocol, host) -> str | None:
        """Return the connection host with the protocol."""
        if not host:
            return host
        if protocol:
            return f"{protocol}://{host}"
        return host

    @staticmethod
    def _normalize_conn_type(conn_type):
        if conn_type == "postgresql":
            conn_type = "postgres"
        elif "-" in conn_type:
            conn_type = conn_type.replace("-", "_")
        return conn_type
