from __future__ import annotations

import base64
import json
import time
from datetime import datetime, timezone
from typing import Any, Optional

import httpx

from src.models.cloud_access import (
    AccessTokenResponse,
    AuthResponse,
    S3Access,
)
from src.utils.config import get_auth_provider, get_auth_service_url, get_github_token
from src.utils.errors import AuthServiceError, NetworkError
from src.utils.logging_utils import get_logger

logger = get_logger("auth")

EXPIRY_BUFFER_SECONDS = 60
AUTH_SERVICE_TIMEOUT = 15.0
AUTH_PATH_ACCESS_TOKEN = "/client/auth"
# Allowed cloud providers forwarded to backend /client/auth.
SUPPORTED_PROVIDERS = {"aws", "alicloud", "gcp"}

_cached_auth_by_provider: dict[str, tuple[int, AuthResponse]] = {}


def decode_jwt_exp(access_token: str) -> Optional[int]:
    try:
        parts = access_token.split(".")
        if len(parts) != 3:
            return None
        payload_b64 = parts[1]
        padding = 4 - len(payload_b64) % 4
        if padding != 4:
            payload_b64 += "=" * padding
        raw = base64.urlsafe_b64decode(payload_b64)
        data = json.loads(raw)
        exp = data.get("exp")
        if exp is None:
            return None
        return int(exp)
    except Exception:
        return None


def _exp_timestamp(response: AuthResponse) -> Optional[int]:
    if response.access_token:
        exp = decode_jwt_exp(response.access_token)
        if exp is not None:
            return exp
    now = int(time.time())
    if response.expires_in is not None and response.expires_in > 0:
        return now + response.expires_in
    return None


def _expires_in_from_expiration(expiration: datetime | None) -> Optional[int]:
    if expiration is None:
        return None
    exp_utc = (
        expiration.replace(tzinfo=timezone.utc)
        if expiration.tzinfo is None
        else expiration.astimezone(timezone.utc)
    )
    remaining = int(exp_utc.timestamp() - time.time())
    return remaining if remaining > 0 else 0


def _normalize_provider(vendor: Optional[str] = None) -> str:
    raw_env_provider = get_auth_provider()
    provider = (vendor or raw_env_provider or "").strip().lower()
    if provider not in SUPPORTED_PROVIDERS:
        supported = ", ".join(sorted(SUPPORTED_PROVIDERS))
        raise AuthServiceError(
            f"provider must be one of: {supported}. "
            f"Got resolved provider: {provider!r} "
            f"(vendor arg={vendor!r}, AUTH_PROVIDER={raw_env_provider!r})"
        )
    return provider


def _set_cached(response: AuthResponse, provider: Optional[str] = None) -> None:
    normalized_provider = _normalize_provider(provider)
    exp = _exp_timestamp(response)
    if exp is None:
        _cached_auth_by_provider.pop(normalized_provider, None)
        return
    _cached_auth_by_provider[normalized_provider] = (exp, response)


def _get_cached(provider: Optional[str] = None) -> Optional[AuthResponse]:
    normalized_provider = _normalize_provider(provider)
    cached = _cached_auth_by_provider.get(normalized_provider)
    if cached is None:
        return None
    exp, response = cached
    if time.time() >= exp - EXPIRY_BUFFER_SECONDS:
        _cached_auth_by_provider.pop(normalized_provider, None)
        return None
    return response


def get_observability_global_service_url(vendor: Optional[str] = None) -> Optional[str]:
    cached = _get_cached(vendor)
    if cached is None:
        return None
    url = cached.observability_global_service
    return url.strip() if url and url.strip() else None


def validate_access_token(access_token: str | None) -> str:
    token = (access_token or "").strip()
    if not token:
        raise AuthServiceError("access_token is required and cannot be empty.")
    exp = decode_jwt_exp(token)
    if exp is not None and time.time() >= exp - EXPIRY_BUFFER_SECONDS:
        raise AuthServiceError("Invalid or expired token.")
    return token


def _ensure_auth_base() -> str:
    base = get_auth_service_url()
    if not base:
        raise AuthServiceError(
            "GitHub auth is not configured. Set AUTH_SERVICE_URL to the backend auth service."
        )
    return base


async def _auth_request(
    method: str,
    path: str,
    bearer_token: Optional[str],
    json_body: Optional[dict[str, Any]] = None,
    params: Optional[dict[str, str]] = None,
) -> dict:
    base = _ensure_auth_base()
    headers = {"Accept": "application/json"}
    if bearer_token:
        headers["Authorization"] = f"Bearer {bearer_token}"
    if json_body is not None:
        headers["Content-Type"] = "application/json"
    url = f"{base}{path}"
    logger.info(
        "auth_request_start method=%s path=%s has_bearer=%s has_json=%s has_params=%s",
        method,
        path,
        bool(bearer_token),
        json_body is not None,
        params is not None,
    )
    try:
        async with httpx.AsyncClient(timeout=AUTH_SERVICE_TIMEOUT) as client:
            response = await client.request(
                method,
                url,
                headers=headers,
                json=json_body,
                params=params,
            )
    except (httpx.NetworkError, httpx.ConnectError, httpx.ConnectTimeout) as e:
        raise NetworkError(
            "Auth service is unreachable. Please verify AUTH_SERVICE_URL and network connectivity."
        ) from e
    except httpx.TimeoutException as e:
        raise AuthServiceError(
            f"Auth service did not respond within {AUTH_SERVICE_TIMEOUT} seconds."
        ) from e

    if response.status_code == 401:
        logger.warning("auth_request_unauthorized path=%s", path)
        raise AuthServiceError("Invalid or expired token.", status_code=401)
    if response.status_code == 403:
        logger.warning("auth_request_forbidden path=%s", path)
        raise AuthServiceError("Access denied by auth service.", status_code=403)
    if response.status_code >= 400:
        logger.warning("auth_request_http_error path=%s status=%s", path, response.status_code)
        raise AuthServiceError(
            f"Auth service error: HTTP {response.status_code}",
            status_code=response.status_code,
        )

    logger.info("auth_request_success path=%s status=%s", path, response.status_code)
    try:
        return response.json()
    except Exception as e:
        raise AuthServiceError("Auth service returned invalid JSON.") from e


async def _fetch_auth(github_token: str, provider: str = "aws") -> AuthResponse:
    """
    Exchange server GitHub token for backend auth payload.

    provider is forwarded to backend /client/auth and defaults to "aws".
    """
    token = (github_token or "").strip()
    if not token:
        raise AuthServiceError("GitHub token is required and cannot be empty.")
    # Normalize provider once so request payload and returned model stay consistent.
    provider_name = (provider or "").strip().lower()
    if provider_name not in SUPPORTED_PROVIDERS:
        supported = ", ".join(sorted(SUPPORTED_PROVIDERS))
        raise AuthServiceError(
            f"provider must be one of: {supported}. Got: {provider!r}"
        )
    data = await _auth_request(
        "POST",
        AUTH_PATH_ACCESS_TOKEN,
        None,
        json_body={"github_token": token, "provider": provider_name},
    )
    try:
        raw = AccessTokenResponse.model_validate(data)
    except Exception as e:
        raise AuthServiceError(f"Auth access_token response invalid: {e}") from e

    resolved_provider = raw.provider or provider_name
    if resolved_provider not in SUPPORTED_PROVIDERS:
        supported = ", ".join(sorted(SUPPORTED_PROVIDERS))
        raise AuthServiceError(
            f"provider in auth response must be one of: {supported}. Got: {resolved_provider!r}"
        )

    # Keep internal cache contract (AuthResponse.expires_in) by deriving seconds from expiration.
    resolved_expires_in = _expires_in_from_expiration(raw.expiration)

    s3_access = None
    if resolved_provider in {"aws", "alicloud"} and raw.access_key_id and raw.secret_access_key:
        s3_access = S3Access(
            provider=resolved_provider,
            access_key_id=raw.access_key_id,
            secret_access_key=raw.secret_access_key,
            region=None,
            endpoint=None,
            account_name=None,
            session_token=raw.session_token,
        )

    return AuthResponse(
        access_token=raw.token,
        gcp_access_token=raw.access_token if resolved_provider == "gcp" else None,
        expires_in=resolved_expires_in,
        auth0Token=raw.auth0_token,
        observabilityGlobalService=raw.observability_global_service,
        s3_access=s3_access,
    )


async def bootstrap_server_auth(vendor: Optional[str] = None) -> AuthResponse:
    logger.info("server_auth_bootstrap_start")
    github_token = get_github_token()
    if not github_token:
        raise AuthServiceError(
            "Server GitHub auth is not configured. Set GITHUB_TOKEN before starting the MCP server."
        )
    provider = _normalize_provider(vendor)
    response = await _fetch_auth(github_token, provider=provider)
    _set_cached(response, provider=provider)
    logger.info("server_auth_bootstrap_complete provider=%s", provider)
    return response


async def get_valid_server_auth(vendor: Optional[str] = None) -> AuthResponse:
    provider = _normalize_provider(vendor)
    cached = _get_cached(provider)
    if cached is not None:
        logger.info("server_auth_cache_hit provider=%s", provider)
        return cached
    logger.info("server_auth_cache_miss provider=%s", provider)
    return await bootstrap_server_auth(provider)


async def get_server_access_token(vendor: Optional[str] = None) -> Optional[str]:
    auth = await get_valid_server_auth(vendor)
    token = auth.access_token
    if not token:
        return None
    return token.strip() or None


async def get_gcp_access_token(vendor: Optional[str] = "gcp") -> Optional[str]:
    auth = await get_valid_server_auth(vendor)
    token = auth.gcp_access_token
    if not token:
        return None
    return token.strip() or None


async def get_server_auth0_token(vendor: Optional[str] = None) -> Optional[str]:
    auth = await get_valid_server_auth(vendor)
    token = auth.auth0_token
    if not token:
        return None
    return token.strip() or None


async def get_delta_credentials(vendor: Optional[str] = None) -> Optional[dict[str, Any]]:
    cached = _get_cached(vendor)
    if cached is None or cached.s3_access is None:
        return None
    logger.info("delta_credentials_found")
    return cached.s3_access.model_dump(exclude_none=True)