import os.path as op
from typing import Any
from urllib.parse import urlsplit

from airflow.configuration import conf
from airflow.plugins_manager import AirflowPlugin
from airflow.security import permissions
from airflow.www.auth import has_access
from airflow.www.views import AirflowBaseView
from flask import abort
from flask_appbuilder import AppBuilder, expose

from cosmos import telemetry
from cosmos.listeners import dag_run_listener, task_instance_listener
from cosmos.plugin.snippets import IFRAME_SCRIPT
from cosmos.plugin.storage import get_storage_type_from_path
from cosmos.settings import dbt_docs_conn_id, dbt_docs_dir, dbt_docs_index_file_name, in_astro_cloud

if in_astro_cloud:
    MENU_ACCESS_PERMISSIONS = [
        (permissions.ACTION_CAN_ACCESS_MENU, "Custom Menu"),
        (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE),
    ]
else:
    MENU_ACCESS_PERMISSIONS = [
        (permissions.ACTION_CAN_READ, permissions.RESOURCE_WEBSITE),
    ]


def bucket_and_key(path: str) -> tuple[str, str]:
    parsed_url = urlsplit(path)
    bucket = parsed_url.netloc
    key = parsed_url.path.lstrip("/")
    return bucket, key


def open_s3_file(path: str, conn_id: str | None) -> str:
    from airflow.providers.amazon.aws.hooks.s3 import S3Hook
    from botocore.exceptions import ClientError

    if conn_id is None:
        conn_id = S3Hook.default_conn_name

    hook = S3Hook(aws_conn_id=conn_id)
    bucket, key = bucket_and_key(path)
    try:
        content = hook.read_key(key=key, bucket_name=bucket)
    except ClientError as e:
        if e.response.get("Error", {}).get("Code", "") == "NoSuchKey":
            raise FileNotFoundError(f"{path} does not exist")
        raise e
    return content  # type: ignore[no-any-return]


def open_gcs_file(path: str, conn_id: str | None) -> str:
    from airflow.providers.google.cloud.hooks.gcs import GCSHook
    from google.cloud.exceptions import NotFound

    if conn_id is None:
        conn_id = GCSHook.default_conn_name

    hook = GCSHook(gcp_conn_id=conn_id)
    bucket, blob = bucket_and_key(path)
    try:
        content = hook.download(bucket_name=bucket, object_name=blob)
    except NotFound:
        raise FileNotFoundError(f"{path} does not exist")
    return content.decode("utf-8")  # type: ignore[no-any-return]


def open_azure_file(path: str, conn_id: str | None) -> str:
    from airflow.providers.microsoft.azure.hooks.wasb import WasbHook
    from azure.core.exceptions import ResourceNotFoundError

    if conn_id is None:
        conn_id = WasbHook.default_conn_name

    hook = WasbHook(wasb_conn_id=conn_id)

    container, blob = bucket_and_key(path)
    try:
        content = hook.read_file(container_name=container, blob_name=blob)
    except ResourceNotFoundError:
        raise FileNotFoundError(f"{path} does not exist")
    return content  # type: ignore[no-any-return]


def open_http_file(path: str, conn_id: str | None) -> str:
    from airflow.providers.http.hooks.http import HttpHook
    from requests.exceptions import HTTPError

    if conn_id is None:
        conn_id = ""

    hook = HttpHook(method="GET", http_conn_id=conn_id)
    try:
        res = hook.run(endpoint=path)
        hook.check_response(res)
    except HTTPError as e:
        if str(e).startswith("404"):
            raise FileNotFoundError(f"{path} does not exist")
        raise e
    return res.text  # type: ignore[no-any-return]


def open_file(path: str, conn_id: str | None = None) -> str:
    """
    Retrieve a file from http, https, gs, s3, or wasb.

    Raise a (base Python) FileNotFoundError if the file is not found.
    """
    if path.strip().startswith("s3://"):
        return open_s3_file(path, conn_id=conn_id)
    elif path.strip().startswith("gs://"):
        return open_gcs_file(path, conn_id=conn_id)
    elif path.strip().startswith("wasb://"):
        return open_azure_file(path, conn_id=conn_id)
    elif path.strip().startswith("http://") or path.strip().startswith("https://"):
        return open_http_file(path, conn_id=conn_id)
    else:
        with open(path) as f:
            content = f.read()
        return content  # type: ignore[no-any-return]


class DbtDocsView(AirflowBaseView):  # type: ignore
    default_view = "dbt_docs"
    route_base = "/cosmos"
    template_folder = op.join(op.dirname(__file__), "templates")
    static_folder = op.join(op.dirname(__file__), "static")

    def create_blueprint(
        self, appbuilder: AppBuilder, endpoint: str | None = None, static_folder: str | None = None
    ) -> None:
        # Make sure the static folder is not overwritten, as we want to use it.
        return super().create_blueprint(appbuilder, endpoint=endpoint, static_folder=self.static_folder)  # type: ignore[no-any-return]

    @expose("/dbt_docs")  # type: ignore[untyped-decorator]
    @has_access(MENU_ACCESS_PERMISSIONS)  # type: ignore[untyped-decorator]
    def dbt_docs(self) -> str:
        # Emit telemetry for dbt docs access
        storage_type = "not_configured"
        if dbt_docs_dir is not None:
            storage_type = get_storage_type_from_path(dbt_docs_dir)

        telemetry.emit_usage_metrics_if_enabled(
            event_type="dbt_docs_access",
            additional_metrics={
                "storage_type": storage_type,
                "docs_dir_configured": dbt_docs_dir is not None,
                "uses_custom_conn": dbt_docs_conn_id is not None,
            },
        )

        if dbt_docs_dir is None:
            return self.render_template("dbt_docs_not_set_up.html")  # type: ignore[no-any-return,no-untyped-call]
        return self.render_template("dbt_docs.html")  # type: ignore[no-any-return,no-untyped-call]

    @expose("/dbt_docs_index.html")  # type: ignore[untyped-decorator]
    @has_access(MENU_ACCESS_PERMISSIONS)  # type: ignore[untyped-decorator]
    def dbt_docs_index(self) -> tuple[str, int, dict[str, Any]]:
        if dbt_docs_dir is None:
            abort(404)
        try:
            html = open_file(op.join(dbt_docs_dir, dbt_docs_index_file_name), conn_id=dbt_docs_conn_id)
        except FileNotFoundError:
            abort(404)
        else:
            html = html.replace("</head>", f"{IFRAME_SCRIPT}</head>")
            return html, 200, {"Content-Security-Policy": "frame-ancestors 'self'"}

    @expose("/catalog.json")  # type: ignore[untyped-decorator]
    @has_access(MENU_ACCESS_PERMISSIONS)  # type: ignore[untyped-decorator]
    def catalog(self) -> tuple[str, int, dict[str, Any]]:
        if dbt_docs_dir is None:
            abort(404)
        try:
            data = open_file(op.join(dbt_docs_dir, "catalog.json"), conn_id=dbt_docs_conn_id)
        except FileNotFoundError:
            abort(404)
        else:
            return data, 200, {"Content-Type": "application/json"}

    @expose("/manifest.json")  # type: ignore[untyped-decorator]
    @has_access(MENU_ACCESS_PERMISSIONS)  # type: ignore[untyped-decorator]
    def manifest(self) -> tuple[str, int, dict[str, Any]]:
        if dbt_docs_dir is None:
            abort(404)
        try:
            data = open_file(op.join(dbt_docs_dir, "manifest.json"), conn_id=dbt_docs_conn_id)
        except FileNotFoundError:
            abort(404)
        else:
            return data, 200, {"Content-Type": "application/json"}


dbt_docs_view = DbtDocsView()


class CosmosPlugin(AirflowPlugin):
    name = "cosmos"
    item = {
        "name": "dbt Docs",
        "category": "Browse",
        "view": dbt_docs_view,
        "href": conf.get("webserver", "base_url") + "/cosmos/dbt_docs",
    }
    appbuilder_views = [item]
    listeners = [dag_run_listener, task_instance_listener]
