from __future__ import annotations

import typing as _t
import mimetypes
import inspect
import pkgutil
import time
import sys
import asyncio
import concurrent.futures
import queue
import threading
from urllib.parse import urlparse

from logging.config import dictConfig
from contextvars import copy_context
from typing import Any, Dict, TYPE_CHECKING

from importlib_metadata import version as _get_distribution_version

# Attempt top-level Quart imports; allow absence if user not using quart backend
try:
    from quart import (
        Quart,
        Response,
        jsonify,
        request,
        Blueprint,
        g as quart_g,
        has_request_context,
        redirect,
        websocket,
    )
except ImportError as _err:
    raise ImportError(
        "All dependencies not installed. Please install it with `dash[quart]` to use the Quart backend."
    ) from _err

import janus

from dash.exceptions import PreventUpdate, InvalidResourceError
from dash.fingerprint import check_fingerprint
from dash._utils import parse_version
from dash import _validate
from .base_server import (
    BaseDashServer,
    RequestAdapter,
    ResponseAdapter,
)
from .ws import (
    DashWebsocketCallback,
    run_ws_sender,
    run_callback_in_executor,
    run_callback_on_loop,
    make_callback_done_handler,
    shutdown_ws_connection,
)
from ._utils import format_traceback_html

if TYPE_CHECKING:
    from dash import Dash


class QuartResponseAdapter(ResponseAdapter):
    """
    A custom Response class that wraps Quart's Response
    and provides a set_response() method for compatibility with Dash's callback system.
    """

    def __init__(self):
        self._quart_response = Response(content_type="application/json")
        super().__init__()

    @property
    def callback_response(self) -> Response:
        return self._quart_response

    def set_cookie(self, key, value="", **kwargs):
        self._quart_response.set_cookie(key, value, **kwargs)

    def append_header(self, key, value):
        self._quart_response.headers.add(key, value)

    def set_header(self, key, value):
        self._quart_response.headers.set(key, value)

    def set_response(self, **kwargs):
        self._quart_response.set_data(kwargs.get("data", ""))
        return self._quart_response


class QuartDashServer(BaseDashServer[Quart]):
    websocket_capability: bool = True

    def __init__(self, server: Quart) -> None:
        super().__init__(server)
        self.server_type = "quart"
        self.config: Dict[str, Any] = {}
        self.error_handling_mode = "ignore"
        self.request_adapter = QuartRequestAdapter
        self.response_adapter = QuartResponseAdapter
        self._active_websockets: set = set()
        self._ws_shutdown_event: asyncio.Event | None = None

    def __call__(self, *args: Any, **kwargs: Any):  # type: ignore[name-defined]
        return self.server(*args, **kwargs)

    @staticmethod
    def create_app(
        name: str = "__main__", config: _t.Optional[_t.Dict[str, _t.Any]] = None
    ):
        if Quart is None:
            raise RuntimeError(
                "Quart is not installed. Install with 'pip install quart' to use the quart backend."
            )
        app = Quart(name)  # type: ignore
        if config:
            for key, value in config.items():
                app.config[key] = value
        return app

    def register_assets_blueprint(
        self, blueprint_name: str, assets_url_path: str, assets_folder: str  # type: ignore[name-defined]
    ):
        bp = Blueprint(
            blueprint_name,
            __name__,
            static_folder=assets_folder,
            static_url_path=assets_url_path,
        )
        self.server.register_blueprint(bp)

    def _get_traceback(self, _secret, error: Exception):
        return format_traceback_html(
            error, self.error_handling_mode, "Quart Debugger", "Quart"
        )

    def register_prune_error_handler(self, secret, prune_errors):
        if prune_errors:
            self.error_handling_mode = "prune"
        else:
            self.error_handling_mode = "raise"

        @self.server.errorhandler(Exception)
        async def _wrap_errors(error):
            if self.error_handling_mode == "ignore":
                return Response(
                    "Internal server error.", status=500, content_type="text/plain"
                )
            tb = self._get_traceback(secret, error)
            return Response(tb, status=500, content_type="text/html")

    def register_timing_hooks(self, _first_run: bool):  # type: ignore[override]  # parity with Flask factory
        @self.server.before_request
        async def _before_request():  # pragma: no cover - timing infra
            if quart_g is not None:
                quart_g.timing_information = {  # type: ignore[attr-defined]
                    "__dash_server": {"dur": time.time(), "desc": None}
                }

        @self.server.after_request
        async def _after_request(response):  # pragma: no cover - timing infra
            timing_information = (
                getattr(quart_g, "timing_information", None)
                if quart_g is not None
                else None
            )
            if timing_information is None:
                return response
            dash_total = timing_information.get("__dash_server", None)
            if dash_total is not None:
                dash_total["dur"] = round((time.time() - dash_total["dur"]) * 1000)
            for name, info in timing_information.items():
                value = name
                if info.get("desc") is not None:
                    value += f';desc="{info["desc"]}"'
                if info.get("dur") is not None:
                    value += f";dur={info['dur']}"
                # Quart/Werkzeug headers expose 'add' (not 'append')
                if hasattr(response.headers, "add"):
                    response.headers.add("Server-Timing", value)
                else:  # fallback just in case
                    response.headers["Server-Timing"] = value
            return response

    def register_error_handlers(self):  # type: ignore[name-defined]
        @self.server.errorhandler(PreventUpdate)
        async def _prevent_update(_):
            return "", 204

        @self.server.errorhandler(InvalidResourceError)
        async def _invalid_resource(err):
            return err.args[0], 404

    def _html_response_wrapper(self, view_func: _t.Callable[..., _t.Any] | str):
        async def wrapped(*_args, **_kwargs):
            html_val = view_func() if callable(view_func) else view_func
            if inspect.iscoroutine(html_val):  # handle async function returning html
                html_val = await html_val
            html = str(html_val)
            return Response(html, content_type="text/html")

        return wrapped

    def add_url_rule(
        self,
        rule: str,
        view_func: _t.Callable[..., _t.Any],
        endpoint: str | None = None,
        methods: list[str] | None = None,
    ):
        self.server.add_url_rule(
            rule, view_func=view_func, endpoint=endpoint, methods=methods or ["GET"]
        )

    def setup_index(self, dash_app: Dash):  # type: ignore[name-defined]
        async def index(*args, **kwargs):
            return Response(dash_app.index(*args, **kwargs), content_type="text/html")  # type: ignore[arg-type]

        # pylint: disable=protected-access
        dash_app._add_url("", index, methods=["GET"])

    def setup_catchall(self, dash_app: Dash):
        async def catchall(
            path: str, *args, **kwargs
        ):  # noqa: ARG001 - path is unused but kept for route signature, pylint: disable=unused-argument
            return Response(dash_app.index(*args, **kwargs), content_type="text/html")  # type: ignore[arg-type]

        # pylint: disable=protected-access
        dash_app._add_url("<path:path>", catchall, methods=["GET"])

    def before_request(self, func: _t.Callable[[], _t.Any]):
        self.server.before_request(func)

    def after_request(self, func: _t.Callable[[], _t.Any]):
        @self.server.after_request
        async def _after(response):
            if func is not None:
                result = func()
                if inspect.iscoroutine(result):  # Allow async hooks
                    await result
            return response

    def has_request_context(self) -> bool:
        if has_request_context is None:
            raise RuntimeError("Quart not installed; cannot check request context")
        return has_request_context()

    # pylint: disable=W0613
    def run(self, dash_app: Dash, host: str, port: int, debug: bool, **kwargs: _t.Any):
        import signal  # pylint: disable=import-outside-toplevel

        # pylint: disable=import-outside-toplevel,import-error
        from hypercorn.config import Config
        from hypercorn.asyncio import serve

        # pylint: enable=import-error

        self.config = {"debug": debug, **kwargs} if debug else kwargs
        # pylint: disable=protected-access
        if dash_app._dev_tools.silence_routes_logging:
            dictConfig(
                {
                    "version": 1,
                    "loggers": {
                        "quart.app": {
                            "level": "ERROR",
                        },
                    },
                }
            )

        # Check if we're running in a non-main thread (e.g., testing context)
        is_main_thread = threading.current_thread() is threading.main_thread()

        config = Config()
        config.bind = [f"{host}:{port}"]
        config.use_reloader = False
        if not is_main_thread:
            config.accesslog = None

        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        # Initialize shutdown event for WebSocket handlers
        self._ws_shutdown_event = asyncio.Event()

        def signal_handler():
            """Handle shutdown signal by setting the WebSocket shutdown event."""
            if self._ws_shutdown_event is not None:
                self._ws_shutdown_event.set()

        # Set up signal handlers in main thread
        if is_main_thread:
            for sig in (signal.SIGINT, signal.SIGTERM):
                try:
                    loop.add_signal_handler(sig, signal_handler)
                except (NotImplementedError, ValueError):
                    pass

            print(f" * Serving Quart app '{self.server.name}'")
            print(f" * Debug mode: {debug}")
            print(
                " * Please use an ASGI server (e.g. Hypercorn) directly in production"
            )
            print(f" * Running on http://{host}:{port} (CTRL + C to quit)")

        async def shutdown_trigger():
            if self._ws_shutdown_event is not None:
                await self._ws_shutdown_event.wait()

        try:
            loop.run_until_complete(
                serve(self.server, config, shutdown_trigger=shutdown_trigger)
            )
        finally:
            loop.close()

    def make_response(
        self,
        data: str | bytes | bytearray,
        mimetype: str | None = None,
        content_type: str | None = None,
        status=None,
    ):
        if Response is None:
            raise RuntimeError("Quart not installed; cannot generate Response")
        return Response(
            data, mimetype=mimetype, content_type=content_type, status=status
        )

    def jsonify(self, obj):
        return jsonify(obj)

    def serve_component_suites(
        self, dash_app: Dash, package_name: str, fingerprinted_path: str
    ):  # noqa: ARG002 unused req preserved for interface parity
        path_in_pkg, has_fingerprint = check_fingerprint(fingerprinted_path)
        _validate.validate_js_path(dash_app.registered_paths, package_name, path_in_pkg)
        extension = "." + path_in_pkg.split(".")[-1]
        mimetype = mimetypes.types_map.get(extension, "application/octet-stream")
        package = sys.modules[package_name]
        dash_app.logger.debug(
            "serving -- package: %s[%s] resource: %s => location: %s",
            package_name,
            getattr(package, "__version__", "unknown"),
            path_in_pkg,
            package.__path__,
        )
        data = pkgutil.get_data(package_name, path_in_pkg)
        headers = {}
        if has_fingerprint:
            headers["Cache-Control"] = "public, max-age=31536000"

        if Response is None:
            raise RuntimeError("Quart not installed; cannot generate Response")
        return Response(data, content_type=mimetype, headers=headers)

    def setup_component_suites(self, dash_app: Dash):
        async def serve(package_name, fingerprinted_path):
            return self.serve_component_suites(
                dash_app, package_name, fingerprinted_path
            )

        # pylint: disable=protected-access
        dash_app._add_url(
            "_dash-component-suites/<string:package_name>/<path:fingerprinted_path>",
            serve,
        )

    def _create_redirect_function(self, redirect_to):
        def _redirect():
            return redirect(redirect_to, code=301)

        return _redirect

    def add_redirect_rule(self, app, fullname, path):
        self.server.add_url_rule(
            fullname,
            fullname,
            self._create_redirect_function(app.get_relative_path(path)),
        )

    # pylint: disable=unused-argument
    def serve_callback(self, dash_app: Dash):  # type: ignore[override]  # Quart always async
        async def _dispatch():
            adapter = QuartRequestAdapter()
            body = await adapter.get_json()
            # pylint: disable=protected-access
            cb_ctx = dash_app._initialize_context(body)
            # pylint: disable=protected-access
            func = dash_app._prepare_callback(cb_ctx, body)
            # pylint: disable=protected-access
            args = dash_app._inputs_to_vals(cb_ctx.inputs_list + cb_ctx.states_list)
            ctx = copy_context()
            # pylint: disable=protected-access
            partial_func = dash_app._execute_callback(
                func, args, cb_ctx.outputs_list, cb_ctx
            )
            response_data = ctx.run(partial_func)
            if inspect.iscoroutine(response_data):  # if user callback is async
                response_data = await response_data
            return cb_ctx.dash_response.set_response(data=response_data)  # type: ignore[arg-type]

        return _dispatch

    def register_callback_api_routes(
        self, callback_api_paths: _t.Dict[str, _t.Callable[..., _t.Any]]
    ):
        """
        Register callback API endpoints on the Quart app.
        Each key in callback_api_paths is a route, each value is a handler (sync or async).
        The view function parses the JSON body and passes it to the handler.
        """
        for path, handler in callback_api_paths.items():
            endpoint = f"dash_callback_api_{path}"
            route = path if path.startswith("/") else f"/{path}"
            methods = ["POST"]

            def _make_view_func(handler):
                if inspect.iscoroutinefunction(handler):

                    async def async_view_func(*args, **kwargs):
                        if request is None:
                            raise RuntimeError(
                                "Quart not installed; request unavailable"
                            )
                        data = await request.get_json()
                        result = await handler(**data) if data else await handler()
                        return jsonify(result)  # type: ignore[arg-type]

                    return async_view_func

                async def sync_view_func(*args, **kwargs):
                    if request is None:
                        raise RuntimeError("Quart not installed; request unavailable")
                    data = await request.get_json()
                    result = handler(**data) if data else handler()
                    return jsonify(result)  # type: ignore[arg-type]

                return sync_view_func

            view_func = _make_view_func(handler)
            self.server.add_url_rule(
                route, endpoint=endpoint, view_func=view_func, methods=methods
            )

    def enable_compression(self) -> None:
        try:
            import quart_compress  # pylint: disable=import-outside-toplevel

            Compress = quart_compress.Compress
            Compress(self.server)
            _flask_compress_version = parse_version(
                _get_distribution_version("quart_compress")
            )
            if not hasattr(
                self.server.config, "COMPRESS_ALGORITHM"
            ) and _flask_compress_version >= parse_version("1.6.0"):
                self.server.config["COMPRESS_ALGORITHM"] = ["gzip"]
        except ImportError as error:
            raise ImportError(
                "To use the compress option, you need to install quart_compress."
            ) from error

    async def _run_ws_hooks(
        self, hooks, ws, *args, default_reason: str = "Rejected"
    ) -> tuple | None:
        """Run WebSocket hooks and return rejection tuple or None if all pass.

        Args:
            hooks: List of hooks to run
            ws: The WebSocket connection
            *args: Additional arguments to pass to hooks
            default_reason: Default reason if hook returns False

        Returns:
            None if all hooks pass, or (code, reason) tuple for rejection
        """
        for hook in hooks:
            try:
                result = hook(ws, *args)
                if inspect.iscoroutine(result):
                    result = await result
                if result is False:
                    return (4001, default_reason)
                if isinstance(result, tuple) and len(result) == 2:
                    return result
            except Exception:  # pylint: disable=broad-exception-caught
                return (4001, "Authentication error")
        return None

    def _validate_ws_origin(
        self, origin: str | None, host: str | None, allowed_origins: list
    ) -> str | None:
        """Validate WebSocket origin. Returns error message or None if valid."""
        if not origin:
            return "Origin header required"
        if origin in allowed_origins:
            return None  # Explicitly allowed
        if not host:
            return "Origin not allowed"
        # Check same-origin
        origin_host = urlparse(origin).netloc
        if origin_host != host:
            return "Origin not allowed"
        return None

    def serve_websocket_callback(self, dash_app: "Dash"):
        """Set up the WebSocket endpoint for callback handling.

        Uses thread pool executor for callback execution with janus queues
        for async/sync communication between main loop and worker threads.

        Args:
            dash_app: The Dash application instance
        """
        # pylint: disable=too-many-statements,too-many-locals
        ws_path = dash_app.config.routes_pathname_prefix + "_dash-ws-callback"
        # pylint: disable=protected-access
        allowed_origins = getattr(dash_app, "_websocket_allowed_origins", [])

        @self.server.websocket(ws_path)
        async def websocket_handler():  # pylint: disable=too-many-branches
            ws = websocket

            # Validate Origin header
            error = self._validate_ws_origin(
                ws.headers.get("origin"), ws.headers.get("host"), allowed_origins
            )
            if error:
                await ws.close(code=4003, reason=error)
                return

            # Call websocket_connect hooks
            # pylint: disable=protected-access
            rejection = await self._run_ws_hooks(
                dash_app._hooks.get_hooks("websocket_connect"),
                ws,
                default_reason="Connection rejected",
            )
            if rejection:
                await ws.close(code=rejection[0], reason=rejection[1])
                return

            await ws.accept()

            # The connection's event loop, used to dispatch async callbacks as
            # tasks and to resolve get_props futures.
            loop = asyncio.get_running_loop()

            # Track this connection for graceful shutdown
            try:
                ws_obj = ws._get_current_object()
                self._active_websockets.add(ws_obj)
            except AttributeError:
                ws_obj = ws
                self._active_websockets.add(ws)

            # Create janus queue for outbound messages (main loop context)
            outbound_queue: janus.Queue[str] = janus.Queue()
            # Track pending get_props requests. Values are queue.Queue (threadpool /
            # sync path) or asyncio.Future (event-loop / async path).
            pending_get_props: Dict[str, queue.Queue | asyncio.Future] = {}
            # Shutdown event to signal connection closure to worker threads
            connection_shutdown_event = threading.Event()
            # Sync callbacks run on a shared thread pool executor (async callbacks
            # run directly on the event loop). A single bounded pool caps the total
            # worker-thread count regardless of how many connections are open.
            executor = self.get_callback_executor(
                getattr(dash_app, "_websocket_max_workers", 4)
            )
            # Track pending callbacks: concurrent.futures.Future (sync/threadpool)
            # or asyncio.Task (async/event-loop).
            pending_callbacks: Dict[
                str, concurrent.futures.Future | asyncio.Future
            ] = {}

            # Start sender task to drain outbound queue (sends pre-serialized text)
            # pylint: disable=protected-access
            batch_delay = getattr(dash_app, "_websocket_batch_delay", 0.005)
            sender_task = asyncio.create_task(
                run_ws_sender(ws.send, outbound_queue, batch_delay)
            )

            try:
                shutdown_event = self._ws_shutdown_event
                while shutdown_event is None or not shutdown_event.is_set():
                    try:
                        # Use timeout to periodically check shutdown event
                        message = await asyncio.wait_for(ws.receive_json(), timeout=1.0)
                    except asyncio.TimeoutError:
                        # Re-check shutdown event (may have been set during run())
                        shutdown_event = self._ws_shutdown_event
                        continue

                    # Call websocket_message hooks
                    rejection = await self._run_ws_hooks(
                        dash_app._hooks.get_hooks("websocket_message"),
                        ws,
                        message,
                        default_reason="Message rejected",
                    )
                    if rejection:
                        await ws.close(code=rejection[0], reason=rejection[1])
                        return

                    msg_type = message.get("type")

                    if msg_type == "callback_request":
                        request_id = message.get("requestId")
                        renderer_id = message.get("rendererId", "")
                        payload = message.get("payload", {})

                        # Validate that the callback is allowed to use WebSocket transport
                        # pylint: disable=protected-access
                        _validate.validate_websocket_callback_request(
                            payload.get("output"),
                            dash_app.callback_map,
                            dash_app._websocket_callbacks,
                        )

                        # Async callbacks (incl. persistent ones) run as tasks on
                        # the event loop; sync callbacks go to the threadpool.
                        # pylint: disable=protected-access
                        cb_spec = dash_app.callback_map.get(payload.get("output"), {})
                        is_async = inspect.iscoroutinefunction(cb_spec.get("callback"))

                        # Create WebSocket callback instance. The loop is passed only
                        # for the async path so get_prop awaits instead of blocking.
                        ws_cb = DashWebsocketCallback(
                            pending_get_props,
                            renderer_id,
                            outbound_queue,
                            connection_shutdown_event,
                            loop if is_async else None,
                        )

                        done_handler = make_callback_done_handler(
                            outbound_queue,
                            pending_callbacks,
                            request_id,
                            renderer_id,
                            connection_shutdown_event,
                        )

                        if is_async:
                            task = asyncio.create_task(
                                run_callback_on_loop(
                                    dash_app,
                                    payload,
                                    ws_cb,
                                    QuartResponseAdapter(),
                                )
                            )
                            task.add_done_callback(done_handler)
                            pending_callbacks[request_id] = task
                        else:
                            # Submit callback to executor
                            future = run_callback_in_executor(
                                executor,
                                dash_app,
                                payload,
                                ws_cb,
                                QuartResponseAdapter(),
                            )
                            # Set up done callback to send response
                            future.add_done_callback(done_handler)
                            pending_callbacks[request_id] = future

                    elif msg_type == "get_props_response":
                        # Resolve the waiting future (async path) or queue (thread
                        # path) for this request (non-blocking).
                        request_id = message.get("requestId")
                        pending = pending_get_props.get(request_id)
                        if isinstance(pending, asyncio.Future):
                            if not pending.done():
                                pending.set_result(message.get("payload"))
                        elif pending is not None:
                            pending.put_nowait(message.get("payload"))

                    elif msg_type == "heartbeat":
                        outbound_queue.sync_q.put_nowait('{"type": "heartbeat_ack"}')

            except asyncio.CancelledError:
                pass  # Server is shutting down, exit gracefully
            except Exception:  # pylint: disable=broad-exception-caught
                pass  # Other exceptions treated as disconnect
            finally:
                self._active_websockets.discard(ws_obj)
                await shutdown_ws_connection(
                    connection_shutdown_event,
                    pending_get_props,
                    pending_callbacks,
                    outbound_queue,
                    sender_task,
                )


class QuartRequestAdapter(RequestAdapter):
    def __init__(self) -> None:
        self._request = request  # type: ignore[assignment]
        if self._request is None:
            raise RuntimeError("Quart not installed; cannot access request context")

    @property
    def context(self):
        if not has_request_context():
            raise RuntimeError("No active request in context")
        return quart_g

    @property
    def request(self) -> _t.Any:
        return self._request

    @property
    def root(self):
        return self.request.root_url

    @property
    def args(self):
        return self.request.args

    @property
    def is_json(self):
        return self.request.is_json

    @property
    def cookies(self):
        return self.request.cookies

    @property
    def headers(self):
        return self.request.headers

    @property
    def full_path(self):
        return self.request.full_path

    @property
    def url(self):
        return str(self.request.url)

    @property
    def remote_addr(self):
        return self.request.remote_addr

    @property
    def origin(self):
        return self.request.headers.get("origin")

    @property
    def path(self):
        return self.request.path

    async def get_json(self):  # pylint: disable=W0236
        # TODO consider using a sync wraper
        return await self.request.get_json()
