import json
import logging
import base64
from typing import Dict, Any, Optional, Tuple, List
import requests

from app.core.config import settings
from app.core.auth.models import UserSession
from app.core.auth.role_resolver import resolve_user_capabilities

logger = logging.getLogger(__name__)


class KeycloakClient:
    def __init__(self):
        self.base_url = settings.KEYCLOAK_INTERNAL_URL.rstrip("/")
        self.realm = settings.KEYCLOAK_REALM
        self.client_id = settings.KEYCLOAK_CLIENT_ID

    @property
    def realm_url(self) -> str:
        return f"{self.base_url}/realms/{self.realm}"

    @property
    def token_endpoint(self) -> str:
        return f"{self.realm_url}/protocol/openid-connect/token"

    @property
    def userinfo_endpoint(self) -> str:
        return f"{self.realm_url}/protocol/openid-connect/userinfo"

    @property
    def jwks_endpoint(self) -> str:
        return f"{self.realm_url}/protocol/openid-connect/certs"

    def exchange_code(self, code: str, redirect_uri: str, code_verifier: Optional[str] = None) -> Dict[str, Any]:
        """Exchange authorization code for tokens from Keycloak."""
        data = {
            "grant_type": "authorization_code",
            "client_id": self.client_id,
            "code": code,
            "redirect_uri": redirect_uri,
        }
        if code_verifier:
            data["code_verifier"] = code_verifier

        resp = requests.post(self.token_endpoint, data=data, timeout=10)
        if not resp.ok:
            logger.error(f"Keycloak code exchange failed: {resp.status_code} - {resp.text}")
            raise ValueError(f"Keycloak token exchange error: {resp.text}")
        return resp.json()

    def parse_token_claims(self, access_token: str) -> Dict[str, Any]:
        """Parse unverified payload claims from Keycloak JWT token."""
        try:
            parts = access_token.split(".")
            if len(parts) != 3:
                raise ValueError("Invalid JWT format")
            payload_b64 = parts[1]
            # Add padding
            rem = len(payload_b64) % 4
            if rem > 0:
                payload_b64 += "=" * (4 - rem)
            payload_json = base64.urlsafe_b64decode(payload_b64).decode("utf-8")
            return json.loads(payload_json)
        except Exception as e:
            logger.error(f"Failed to parse Keycloak access token: {e}")
            raise ValueError(f"Invalid access token: {e}")

    def build_user_session_from_token(self, token_data: Dict[str, Any]) -> UserSession:
        """Construct unified UserSession from Keycloak access_token and id_token."""
        access_token = token_data.get("access_token", "")
        claims = self.parse_token_claims(access_token)

        sub = claims.get("sub", "")
        username = claims.get("preferred_username") or claims.get("username") or sub
        name = claims.get("name") or username
        email = claims.get("email")

        # Extract roles from realm_access or direct roles claim
        raw_roles: List[str] = []
        if isinstance(claims.get("roles"), list):
            raw_roles.extend(claims["roles"])
        realm_access = claims.get("realm_access", {})
        if isinstance(realm_access.get("roles"), list):
            raw_roles.extend(realm_access["roles"])

        # Extract groups
        raw_groups: List[str] = claims.get("groups", [])
        if not isinstance(raw_groups, list):
            raw_groups = []

        # Resolve internal capabilities dynamically
        cap, can_edit, can_deploy, is_admin, is_readonly = resolve_user_capabilities(raw_roles, raw_groups)

        return UserSession(
            id=sub,
            username=username,
            name=name,
            email=email,
            auth_provider="keycloak",
            raw_roles=raw_roles,
            raw_groups=raw_groups,
            capability=cap,
            can_edit=can_edit,
            can_deploy=can_deploy,
            is_admin=is_admin,
            is_readonly=is_readonly
        )


keycloak_client = KeycloakClient()