# 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.
import logging
from typing import Any

import backoff
from flask_appbuilder.api import expose, protect, request, rison, safe
from flask_appbuilder.models.sqla.interface import SQLAInterface

from superset import db, event_logger
from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
from superset.daos.query import QueryDAO
from superset.databases.filters import DatabaseFilter
from superset.exceptions import SupersetException
from superset.models.sql_lab import Query
from superset.queries.filters import QueryFilter
from superset.queries.schemas import (
    openapi_spec_methods_override,
    queries_get_updated_since_schema,
    QuerySchema,
    StopQuerySchema,
)
from superset.superset_typing import FlaskResponse
from superset.views.base_api import (
    BaseSupersetModelRestApi,
    RelatedFieldFilter,
    requires_json,
    statsd_metrics,
)
from superset.views.filters import BaseFilterRelatedUsers, FilterRelatedOwners

logger = logging.getLogger(__name__)


class QueryRestApi(BaseSupersetModelRestApi):
    datamodel = SQLAInterface(Query)

    resource_name = "query"

    class_permission_name = "Query"
    method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP

    allow_browser_login = True
    include_route_methods = {
        RouteMethod.GET,
        RouteMethod.GET_LIST,
        RouteMethod.RELATED,
        RouteMethod.DISTINCT,
        "stop_query",
        "get_updated_since",
    }

    apispec_parameter_schemas = {
        "queries_get_updated_since_schema": queries_get_updated_since_schema,
    }

    list_columns = [
        "id",
        "changed_on",
        "client_id",
        "database.id",
        "database.database_name",
        "executed_sql",
        "error_message",
        "limit",
        "limiting_factor",
        "progress",
        "rows",
        "schema",
        "select_as_cta",
        "sql",
        "sql_editor_id",
        "sql_tables",
        "status",
        "tab_name",
        "user.first_name",
        "user.id",
        "user.last_name",
        "start_time",
        "end_time",
        "tmp_table_name",
        "tracking_url",
        "results_key",
    ]
    show_columns = [
        "id",
        "changed_on",
        "client_id",
        "database.id",
        "end_result_backend_time",
        "end_time",
        "error_message",
        "executed_sql",
        "limit",
        "progress",
        "results_key",
        "rows",
        "schema",
        "select_as_cta",
        "select_as_cta_used",
        "select_sql",
        "sql",
        "sql_editor_id",
        "start_running_time",
        "start_time",
        "status",
        "tab_name",
        "tmp_schema_name",
        "tmp_table_name",
        "tracking_url",
    ]
    base_filters = [["id", QueryFilter, lambda: []]]
    base_order = ("changed_on", "desc")
    list_model_schema = QuerySchema()
    stop_query_schema = StopQuerySchema()

    openapi_spec_tag = "Queries"
    openapi_spec_methods = openapi_spec_methods_override
    openapi_spec_component_schemas = (StopQuerySchema,)

    order_columns = [
        "changed_on",
        "database.database_name",
        "rows",
        "schema",
        "start_time",
        "sql",
        "tab_name",
        "user.first_name",
    ]
    base_related_field_filters = {
        "created_by": [["id", BaseFilterRelatedUsers, lambda: []]],
        "changed_by": [["id", BaseFilterRelatedUsers, lambda: []]],
        "user": [["id", BaseFilterRelatedUsers, lambda: []]],
        "database": [["id", DatabaseFilter, lambda: []]],
    }
    related_field_filters = {
        "created_by": RelatedFieldFilter("first_name", FilterRelatedOwners),
        "changed_by": RelatedFieldFilter("first_name", FilterRelatedOwners),
        "user": RelatedFieldFilter("first_name", FilterRelatedOwners),
    }

    search_columns = [
        "changed_on",
        "database",
        "sql",
        "status",
        "user",
        "start_time",
        "sql_editor_id",
    ]

    allowed_rel_fields = {"database", "user"}
    allowed_distinct_fields = {"status"}

    @expose("/updated_since")
    @protect()
    @safe
    @rison(queries_get_updated_since_schema)
    @statsd_metrics
    @event_logger.log_this_with_context(
        action=lambda self, *args, **kwargs: f"{self.__class__.__name__}"
        f".get_updated_since",
        log_to_statsd=False,
    )
    def get_updated_since(self, **kwargs: Any) -> FlaskResponse:
        """Get a list of queries that changed after last_updated_ms.
        ---
        get:
          summary: Get a list of queries that changed after last_updated_ms
          parameters:
          - in: query
            name: q
            content:
              application/json:
                schema:
                  $ref: '#/components/schemas/queries_get_updated_since_schema'
          responses:
            200:
              description: Queries list
              content:
                application/json:
                  schema:
                    type: object
                    properties:
                      result:
                        description: >-
                          A List of queries that changed after last_updated_ms
                        type: array
                        items:
                          $ref: '#/components/schemas/{{self.__class__.__name__}}.get'
            400:
              $ref: '#/components/responses/400'
            401:
              $ref: '#/components/responses/401'
            404:
              $ref: '#/components/responses/404'
            500:
              $ref: '#/components/responses/500'
        """
        try:
            last_updated_ms = kwargs["rison"].get("last_updated_ms", 0)
            queries = QueryDAO.get_queries_changed_after(last_updated_ms)
            payload = [q.to_dict() for q in queries]
            return self.response(200, result=payload)
        except SupersetException as ex:
            return self.response(ex.status, message=ex.message)

    @expose("/stop", methods=("POST",))
    @protect()
    @safe
    @statsd_metrics
    @event_logger.log_this_with_context(
        action=lambda self, *args, **kwargs: f"{self.__class__.__name__}"
        f".stop_query",
        log_to_statsd=False,
    )
    @backoff.on_exception(
        backoff.constant,
        Exception,
        interval=1,
        on_backoff=lambda details: db.session.rollback(),  # pylint: disable=consider-using-transaction
        on_giveup=lambda details: db.session.rollback(),  # pylint: disable=consider-using-transaction
        max_tries=5,
    )
    @requires_json
    def stop_query(self) -> FlaskResponse:
        """Manually stop a query with client_id.
        ---
        post:
          summary: Manually stop a query with client_id
          requestBody:
            description: Stop query schema
            required: true
            content:
              application/json:
                schema:
                  $ref: '#/components/schemas/StopQuerySchema'
          responses:
            200:
              description: Query stopped
              content:
                application/json:
                  schema:
                    type: object
                    properties:
                        result:
                            type: string
            400:
              $ref: '#/components/responses/400'
            401:
              $ref: '#/components/responses/401'
            404:
              $ref: '#/components/responses/404'
            500:
              $ref: '#/components/responses/500'
        """
        try:
            body = self.stop_query_schema.load(request.json)
            QueryDAO.stop_query(body["client_id"])
            return self.response(200, result="OK")
        except SupersetException as ex:
            return self.response(ex.status, message=ex.message)
