from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from great_expectations.execution_engine import PandasExecutionEngine
from great_expectations.expectations.expectation import ColumnMapExpectation
from great_expectations.expectations.metrics import (
    ColumnMapMetricProvider,
    column_condition_partial,
)

BASE_API_URL = "https://content.airhex.com/api/v3.7.4/airports/?iata="


# This class defines a Metric to support your Expectation.
# For most ColumnMapExpectations, the main business logic for calculation will live in this class.
class ColumnValuesToBeValidIataCode(ColumnMapMetricProvider):
    # This is the id string that will be used to reference your metric.
    condition_metric_name = "column_values.valid_iata_code"

    # This method implements the core logic for the PandasExecutionEngine
    @column_condition_partial(engine=PandasExecutionEngine)
    def _pandas(cls, column, **kwargs):
        def valid_iata_code(x):
            if not isinstance(x, str):
                return False
            if len(x) != 3:
                return False
            try:
                urlopen(
                    Request(
                        BASE_API_URL + x, headers={"User-Agent": "great_expectations"}
                    )
                )
                return True
            except (URLError, HTTPError):
                return False

        return column.apply(lambda x: valid_iata_code(x) if x else False)

    # This method defines the business logic for evaluating your metric when using a SqlAlchemyExecutionEngine
    # @column_condition_partial(engine=SqlAlchemyExecutionEngine)
    # def _sqlalchemy(cls, column, _dialect, **kwargs):
    #     raise NotImplementedError

    # This method defines the business logic for evaluating your metric when using a SparkDFExecutionEngine
    # @column_condition_partial(engine=SparkDFExecutionEngine)
    # def _spark(cls, column, **kwargs):
    #     raise NotImplementedError


# This class defines the Expectation itself
class ExpectColumnValuesToBeValidIataCode(ColumnMapExpectation):
    """Expect column values to be valid IATA airport codes."""

    # These examples will be shown in the public gallery.
    # They will also be executed as unit tests for your Expectation.
    examples = [
        {
            # https://www.ccra.com/airport-codes has a searchable list of codes
            "data": {
                "valid_airport_codes": ["MCO", "SFO", "LAX", "JFK", "LGA", "EWR"],
                "invalid_airport_codes": [
                    "XYZ",
                    "UWU",
                    "MEME",
                    "FFF",
                    "ZZZ",
                    "ABC",
                ],
            },
            "tests": [
                {
                    "title": "basic_positive_test",
                    "exact_match_out": False,
                    "include_in_gallery": True,
                    "in": {"column": "valid_airport_codes"},
                    "out": {"success": True},
                },
                {
                    "title": "basic_negative_test",
                    "exact_match_out": False,
                    "include_in_gallery": True,
                    "in": {"column": "invalid_airport_codes"},
                    "out": {"success": False},
                },
            ],
        }
    ]

    # This is the id string of the Metric used by this Expectation.
    # For most Expectations, it will be the same as the `condition_metric_name` defined in your Metric class above.
    map_metric = "column_values.valid_iata_code"

    # This is a list of parameter names that can affect whether the Expectation evaluates to True or False
    success_keys = ("mostly",)

    # This dictionary contains default values for any parameters that should have default values
    default_kwarg_values = {}

    # This object contains metadata for display in the public Gallery
    library_metadata = {
        "maturity": "experimental",
        "tags": ["experimental", "expectation", "validation"],
        "contributors": [
            "@rdodev",
        ],
    }


if __name__ == "__main__":
    ExpectColumnValuesToBeValidIataCode().print_diagnostic_checklist()
