from __future__ import annotations

import asyncio
from typing import Any

import httpx

from mcp_eregistrations_keycloak.keycloak_client.errors import (
    KeycloakAuthError,
    KeycloakClientError,
    KeycloakConflictError,
    KeycloakForbiddenError,
    KeycloakNotFoundError,
    KeycloakServerError,
    KeycloakValidationError,
)

_DEFAULT_TIMEOUT = 30.0
_MAX_RESPONSE_SIZE = 100_000  # 100KB — truncate larger responses
_RETRYABLE_STATUSES = {429, 502, 503, 504}
_IDEMPOTENT_METHODS = {"GET", "PUT", "HEAD", "OPTIONS"}
_MAX_RETRIES = 3


class KeycloakAdminClient:
    """Async HTTP client for the Keycloak Admin REST API.

    Provides two access modes:
    - Admin API (base_url + /admin): realm/user/role management
    - Public endpoints (base_url): OIDC token, JWKS, introspection
    """

    def __init__(
        self,
        base_url: str,
        realm: str,
        instance: str | None = None,
        timeout: float = _DEFAULT_TIMEOUT,
    ):
        self.base_url = base_url.rstrip("/")
        self.default_realm = realm
        self._instance = instance
        self._timeout = timeout
        self._token: str | None = None
        self._client: httpx.AsyncClient | None = None
        self._refresh_lock = asyncio.Lock()

    @property
    def public_base_url(self) -> str:
        """Base URL without /admin — for OIDC/token/JWKS endpoints."""
        return self.base_url

    @classmethod
    def for_instance(cls, instance: str | None = None) -> KeycloakAdminClient:
        """Factory: resolves instance config, returns context manager.

        Raises ToolError if instance uses CAS (no Keycloak config).
        """
        from fastmcp.exceptions import ToolError

        from mcp_eregistrations_common.config import get_instance_config

        config = get_instance_config(instance)

        keycloak_url = config.get("keycloak_url")
        if not keycloak_url:
            raise ToolError(
                f"Instance '{instance}' uses CAS authentication and has no "
                f"Keycloak Admin API. Keycloak tools only work with "
                f"Keycloak-authenticated instances."
            )

        return cls(
            base_url=keycloak_url,
            realm=config.get("keycloak_realm", "master"),
            instance=instance,
        )

    async def _resolve_token(self) -> str | None:
        if self._instance:
            from mcp_eregistrations_common.auth import get_or_refresh_token

            return await get_or_refresh_token(self._instance)
        return self._token

    async def __aenter__(self) -> KeycloakAdminClient:
        self._token = await self._resolve_token()
        if not self._token:
            from fastmcp.exceptions import ToolError

            raise ToolError(
                "Not authenticated. Run kc_auth_login first to store credentials."
            )
        self._client = httpx.AsyncClient(
            base_url=f"{self.base_url}/admin",
            headers={"Authorization": f"Bearer {self._token}"},
            timeout=self._timeout,
        )
        return self

    async def __aexit__(self, *args: Any) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None

    def _ensure_client(self) -> httpx.AsyncClient:
        assert self._client is not None, (
            "Use 'async with KeycloakAdminClient(...)' first"
        )
        return self._client

    async def _refresh_and_update_header(self) -> bool:
        """Refresh token with mutex to prevent concurrent refreshes."""
        async with self._refresh_lock:
            new_token = await self._resolve_token()
            if new_token and new_token != self._token:
                self._token = new_token
                self._ensure_client().headers["Authorization"] = f"Bearer {new_token}"
                return True
            return False

    async def _request_with_retry(
        self,
        method: str,
        path: str,
        *,
        params: dict[str, Any] | None = None,
        json: dict[str, Any] | list[Any] | None = None,
    ) -> Any:
        """Execute request with 401 retry + exponential backoff for transient errors."""
        client = self._ensure_client()
        resp = await client.request(method, path, params=params, json=json)

        # 401: try token refresh once
        if resp.status_code == 401 and self._instance:
            refreshed = await self._refresh_and_update_header()
            if refreshed:
                resp = await client.request(method, path, params=params, json=json)

        # Transient errors: exponential backoff (idempotent methods only)
        if resp.status_code in _RETRYABLE_STATUSES and method in _IDEMPOTENT_METHODS:
            for attempt in range(1, _MAX_RETRIES + 1):
                delay = min(2**attempt, 10)  # 2s, 4s, 8s (capped at 10s)
                await asyncio.sleep(delay)
                resp = await client.request(method, path, params=params, json=json)
                if resp.status_code not in _RETRYABLE_STATUSES:
                    break

        return self._handle_response(resp)

    # --- Admin API methods (routed through /admin base) ---

    async def get(self, path: str, **params: Any) -> Any:
        return await self._request_with_retry("GET", path, params=params or None)

    async def post(
        self, path: str, data: dict[str, Any] | list[Any] | None = None
    ) -> Any:
        return await self._request_with_retry("POST", path, json=data)

    async def put(self, path: str, data: dict[str, Any] | None = None) -> Any:
        return await self._request_with_retry("PUT", path, json=data)

    async def delete(self, path: str, data: list[Any] | None = None) -> Any:
        """DELETE with optional JSON body (used by Keycloak for role removal)."""
        return await self._request_with_retry("DELETE", path, json=data)

    # --- Public endpoint methods (bypass /admin base) ---

    async def post_form(self, url: str, data: dict[str, str]) -> Any:
        """POST with form-urlencoded body to a full URL (not under /admin).

        Used for OIDC token endpoint, introspection, etc.
        Does NOT add Bearer auth header (caller provides credentials in form data).
        """
        async with httpx.AsyncClient(timeout=self._timeout) as public_client:
            resp = await public_client.post(
                url,
                data=data,
                headers={"Content-Type": "application/x-www-form-urlencoded"},
            )
            return self._handle_response(resp)

    async def public_get(self, url: str) -> Any:
        """GET from a full URL (not under /admin).

        Used for JWKS endpoint, OIDC discovery, etc.
        """
        async with httpx.AsyncClient(timeout=self._timeout) as public_client:
            resp = await public_client.get(url)
            return self._handle_response(resp)

    # --- Response handling ---

    def _handle_response(self, resp: httpx.Response) -> Any:
        """Check response status and return JSON or raise appropriate error."""
        if resp.status_code == 204:
            return None

        if resp.status_code == 201:
            location = resp.headers.get("Location")
            if location:
                return {"id": location.rsplit("/", 1)[-1]}
            return {}

        if resp.status_code == 200:
            if not resp.content:
                return {}
            # Large response guard
            if len(resp.content) > _MAX_RESPONSE_SIZE:
                data = resp.json()
                if isinstance(data, list) and len(data) > 100:
                    return {
                        "items": data[:100],
                        "total": len(data),
                        "truncated": True,
                        "message": (
                            f"Response truncated: {len(data)} items, showing first 100."
                        ),
                    }
                # Non-list large responses returned as-is (realm config, etc.)
            return resp.json()

        # Error responses — extract message, strip internals
        try:
            body = resp.json()
        except Exception:
            body = resp.text
        msg = str(
            body.get("errorMessage", body.get("error", str(body)))
            if isinstance(body, dict)
            else body
        )

        if resp.status_code == 401:
            raise KeycloakAuthError(msg, status_code=401)
        if resp.status_code == 403:
            raise KeycloakForbiddenError(msg, status_code=403)
        if resp.status_code == 404:
            raise KeycloakNotFoundError(msg, status_code=404)
        if resp.status_code == 409:
            raise KeycloakConflictError(msg, status_code=409)
        if resp.status_code == 400:
            raise KeycloakValidationError(msg, status_code=400)
        if resp.status_code >= 500:
            raise KeycloakServerError(msg, status_code=resp.status_code)

        raise KeycloakClientError(msg, status_code=resp.status_code)
