from __future__ import annotations

from dataclasses import replace
from typing import Any, TYPE_CHECKING
from urllib.parse import urlparse, urlsplit, urlunsplit

from relationalai.runtime.execution import Executor, HttpRequest
from relationalai.services._shared.da_request import DA_BASE_URL_PROVIDER_META_KEY

if TYPE_CHECKING:
    from relationalai.auth import TokenHandler


def _replace_url_base(url: str, new_base_url: str) -> str:
    src = urlsplit(url)
    dst = urlsplit(new_base_url)
    return urlunsplit((dst.scheme, dst.netloc, src.path, src.query, src.fragment))


class DirectAccessAuthMiddleware:
    """Inject upstream-style Direct Access auth headers per request.

    - `Authorization: Snowflake Token=\"...\"`
    - `X-SF-SPCS-Authentication-Method: OAUTH`

    Tokens are minted/refreshed via the provided `TokenHandler`.
    """

    def __init__(self, *, token_handler: "TokenHandler | None" = None, static_token: str | None = None):
        if token_handler is None and not static_token:
            raise ValueError("DirectAccessAuthMiddleware requires token_handler or static_token.")
        if token_handler is not None and static_token:
            raise ValueError("DirectAccessAuthMiddleware accepts either token_handler or static_token, not both.")
        self._handler = token_handler
        self._static_token = str(static_token) if static_token else None

    def wrap(self, nxt: Executor[HttpRequest, Any]) -> Executor[HttpRequest, Any]:
        handler = self._handler
        static_token = self._static_token

        class _Wrapped(Executor[HttpRequest, Any]):
            def execute(_self: object, req: HttpRequest) -> Any:
                if getattr(req, "kind", None) != "http":
                    return nxt.execute(req)

                def _service_endpoint(url: str) -> str:
                    p = urlparse(url)
                    if p.scheme and p.netloc:
                        return f"{p.scheme}://{p.netloc}"
                    return url

                def _inject(url: str) -> HttpRequest:
                    if static_token is not None:
                        token = static_token
                    else:
                        assert handler is not None
                        token = handler.get_ingress_token(_service_endpoint(url))
                    hdrs: dict[str, str] = dict(req.headers or {})
                    hdrs["Authorization"] = f'Snowflake Token="{token}"'
                    hdrs.setdefault("X-SF-SPCS-Authentication-Method", "OAUTH")
                    return replace(req, url=url, headers=hdrs)

                try:
                    return nxt.execute(_inject(req.url))
                except Exception as e:
                    if static_token is not None or handler is None:
                        raise
                    from relationalai.auth import DirectAccessEndpointStaleError

                    if not isinstance(e, DirectAccessEndpointStaleError):
                        raise
                    provider = (req.meta or {}).get(DA_BASE_URL_PROVIDER_META_KEY)
                    if not callable(provider):
                        raise
                    refreshed_url = _replace_url_base(req.url, str(provider(True)))
                    return nxt.execute(_inject(refreshed_url))

        return _Wrapped()


__all__ = ["DirectAccessAuthMiddleware"]

