from __future__ import annotations

import base64
import json
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from typing import Any

from .config import Config


@dataclass(frozen=True)
class OidcProvider:
    authorization_endpoint: str
    token_endpoint: str
    userinfo_endpoint: str
    end_session_endpoint: str | None = None


def _json_get(url: str) -> dict[str, Any]:
    with urllib.request.urlopen(url, timeout=20) as response:
        return json.loads(response.read().decode())


def discover(config: Config) -> OidcProvider:
    metadata = _json_get(f"{config.oidc_issuer}/.well-known/openid-configuration")
    return OidcProvider(
        authorization_endpoint=metadata["authorization_endpoint"],
        token_endpoint=metadata["token_endpoint"],
        userinfo_endpoint=metadata["userinfo_endpoint"],
        end_session_endpoint=metadata.get("end_session_endpoint"),
    )


def authorization_url(config: Config, state: str, next_path: str) -> str:
    provider = discover(config)
    query = urllib.parse.urlencode(
        {
            "client_id": config.oidc_client_id,
            "redirect_uri": f"{config.base_url}/auth/callback",
            "response_type": "code",
            "scope": config.oidc_scopes,
            "state": state,
        }
    )
    return f"{provider.authorization_endpoint}?{query}"


def exchange_code(config: Config, code: str) -> dict[str, Any]:
    provider = discover(config)
    body = urllib.parse.urlencode(
        {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": f"{config.base_url}/auth/callback",
        }
    ).encode()
    request = urllib.request.Request(provider.token_endpoint, data=body, method="POST")
    auth = base64.b64encode(f"{config.oidc_client_id}:{config.oidc_client_secret}".encode()).decode()
    request.add_header("Authorization", f"Basic {auth}")
    request.add_header("Content-Type", "application/x-www-form-urlencoded")
    try:
        with urllib.request.urlopen(request, timeout=20) as response:
            return json.loads(response.read().decode())
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="replace")
        raise RuntimeError(f"OIDC token exchange failed: HTTP {exc.code} {detail}") from exc


def userinfo(config: Config, access_token: str) -> dict[str, Any]:
    provider = discover(config)
    request = urllib.request.Request(provider.userinfo_endpoint)
    request.add_header("Authorization", f"Bearer {access_token}")
    try:
        with urllib.request.urlopen(request, timeout=20) as response:
            return json.loads(response.read().decode())
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode(errors="replace")
        raise RuntimeError(f"OIDC userinfo failed: HTTP {exc.code} {detail}") from exc


def groups_from_claims(claims: dict[str, Any]) -> set[str]:
    raw = claims.get("groups") or claims.get("ak_groups") or []
    if isinstance(raw, str):
        return {raw}
    if isinstance(raw, list):
        return {str(group) for group in raw}
    return set()


def logout_url(config: Config) -> str:
    provider = discover(config)
    # Authentik advertises {issuer}/end-session/; fall back to that shape if discovery omits it.
    endpoint = provider.end_session_endpoint or f"{config.oidc_issuer}/end-session/"
    query = urllib.parse.urlencode({"post_logout_redirect_uri": config.base_url})
    return f"{endpoint}?{query}"
