# SPDX-License-Identifier: Apache-2.0.
import logging
from typing import List, Optional, TYPE_CHECKING, Dict, Tuple
from urllib.parse import urlparse

from openlineage.airflow.extractors.dbapi_utils import (
    get_table_schemas,
    create_information_schema_query,
    TablesHierarchy
)
from openlineage.airflow.utils import get_connection
from openlineage.airflow.extractors.base import BaseExtractor, TaskMetadata
from openlineage.client.facet import BaseFacet, SqlJobFacet
from openlineage.common.sql import SqlMeta, parse
from openlineage.common.dataset import Dataset, Source
from abc import abstractmethod

from openlineage.common.sql.parser import DbTableMeta

if TYPE_CHECKING:
    from airflow.models import Connection
    from airflow.hooks.base import BaseHook

logger = logging.getLogger(__name__)


class SqlExtractor(BaseExtractor):
    default_schema = "public"

    _information_schema_columns = [
        "table_schema",
        "table_name",
        "column_name",
        "ordinal_position",
        "udt_name",
    ]
    _information_schema_table_name = "information_schema.columns"
    _is_information_schema_cross_db = False
    _is_case_sensitive = False

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._conn = None
        self._hook = None
        self._database: Optional[str] = None
        self._scheme: Optional[str] = None

    def extract(self) -> TaskMetadata:
        task_name = f"{self.operator.dag_id}.{self.operator.task_id}"
        job_facets = {"sql": SqlJobFacet(self.operator.sql)}
        run_facets: Dict = {}

        # (1) Parse sql statement to obtain input / output tables.
        logger.debug(f"Sending SQL to parser: {self.operator.sql}")
        sql_meta: Optional[SqlMeta] = parse(self.operator.sql, self.default_schema)
        logger.debug(f"Got meta {sql_meta}")

        if not sql_meta:
            return TaskMetadata(
                name=task_name,
                inputs=[],
                outputs=[],
                run_facets=run_facets,
                job_facets=job_facets,
            )

        # (2) Construct source object
        source = Source(
            scheme=self.scheme,
            authority=self._get_authority(),
            connection_url=self._get_connection_uri(),
        )

        database = getattr(self.operator, "database", None)
        if not database:
            database = self._get_database()

        # (3) Map input / output tables to dataset objects with source set
        # as the current connection. We need to also fetch the schema for the
        # input tables to format the dataset name as:
        # {schema_name}.{table_name}
        inputs, outputs = get_table_schemas(
            self.hook,
            source,
            database,
            self._information_schema_query(sql_meta.in_tables)
            if sql_meta.in_tables
            else None,
            self._information_schema_query(sql_meta.out_tables)
            if sql_meta.out_tables
            else None,
        )

        for ds in inputs:
            ds.input_facets = self._get_input_facets()

        for ds in outputs:
            ds.output_facets = self._get_output_facets()

        db_specific_run_facets = self._get_db_specific_run_facets(
            source, inputs, outputs
        )

        run_facets = {**run_facets, **db_specific_run_facets}

        return TaskMetadata(
            name=task_name,
            inputs=[ds.to_openlineage_dataset() for ds in inputs],
            outputs=[ds.to_openlineage_dataset() for ds in outputs],
            run_facets=run_facets,
            job_facets=job_facets,
        )

    def _conn_id(self) -> str:
        return getattr(self.hook, self.hook.conn_name_attr)

    @property
    def hook(self):
        if not self._hook:
            self._hook = self._get_hook()
        return self._hook

    @property
    def conn(self) -> "Connection":
        if not self._conn:
            self._conn = get_connection(self._conn_id())
        return self._conn  # type: ignore

    @property
    def scheme(self) -> Optional[str]:
        if not self._scheme:
            self._scheme = self._get_scheme()
        return self._scheme

    @property
    def database(self) -> Optional[str]:
        if not self._database:
            self._database = self._get_database()
        return self._database

    @abstractmethod
    def _get_scheme(self) -> Optional[str]:
        raise NotImplementedError

    @abstractmethod
    def _get_database(self) -> Optional[str]:
        raise NotImplementedError

    def _get_authority(self) -> str:
        if self.conn.host and self.conn.port:
            return f"{self.conn.host}:{self.conn.port}"
        else:
            parsed = urlparse(self.conn.get_uri())
            return f"{parsed.hostname}:{parsed.port}"

    @abstractmethod
    def _get_hook(self) -> "BaseHook":
        raise NotImplementedError

    @abstractmethod
    def _get_connection_uri(self) -> str:
        raise NotImplementedError

    def _get_db_specific_run_facets(
        self,
        source: Source,
        inputs: Tuple[List[Dataset], ...],
        outputs: Tuple[List[Dataset], ...],
    ) -> Dict[str, BaseFacet]:
        return {}

    def _get_input_facets(self) -> Dict[str, BaseFacet]:
        return {}

    def _get_output_facets(self) -> Dict[str, BaseFacet]:
        return {}

    def _information_schema_query(self, tables: List[DbTableMeta]) -> str:
        tables_hierarchy = self._get_tables_hierarchy(
            tables,
            database=self.database,
            is_cross_db=self._is_information_schema_cross_db,
        )
        return create_information_schema_query(
            columns=self._information_schema_columns,
            information_schema_table_name=self._information_schema_table_name,
            tables_hierarchy=tables_hierarchy,
            case_sensitive=self._is_case_sensitive,
        )

    @staticmethod
    def _get_tables_hierarchy(
        tables: List[DbTableMeta],
        database: Optional[str] = None,
        is_cross_db: bool = False,
    ) -> TablesHierarchy:
        def normalize_name(name: Optional[str]) -> Optional[str]:
            return name.lower() if name else name
        hierarchy: TablesHierarchy = {}
        for table in tables:
            if is_cross_db:
                db = table.database or database
            else:
                db = None
            hierarchy.setdefault(normalize_name(db), {}).setdefault(
                normalize_name(table.schema), []
            ).append(table.name)
        return hierarchy
