from __future__ import annotations

import base64
import hashlib
import uuid
from datetime import UTC
from datetime import datetime
from datetime import timedelta
from html import escape
from importlib import resources
from urllib.parse import urlparse

import jwt
from jwt.types import Options

from amsdal_ml.mcp_server.mcp_oauth.services.models import JwtDecodedPayload
from amsdal_ml.mcp_server.mcp_oauth.services.types import DEFAULT_NEXT_PATH
from amsdal_ml.mcp_server.mcp_oauth.services.types import JWT_ALGORITHM
from amsdal_ml.mcp_server.mcp_oauth.services.types import LOGIN_TEMPLATE_PATH
from amsdal_ml.mcp_server.mcp_oauth.services.types import TEXT_ENCODING_UTF8
from amsdal_ml.mcp_server.mcp_oauth.services.types import AuthorizationCodeCacheEntry
from amsdal_ml.mcp_server.mcp_oauth.services.types import JwtClaims
from amsdal_ml.mcp_server.mcp_oauth.services.types import OAuthConstants
from amsdal_ml.ml_config import ml_config

_CODE_JTI_CACHE: dict[str, AuthorizationCodeCacheEntry] = {}


def _cleanup_cache() -> None:
    now_ts = int(datetime.now(tz=UTC).timestamp())
    to_del = [j for j, entry in _CODE_JTI_CACHE.items() if entry['exp_ts'] <= now_ts]
    for j in to_del:
        _CODE_JTI_CACHE.pop(j, None)


def store_authorization_code_jti(jti: str, exp_ts: int) -> None:
    _cleanup_cache()
    _CODE_JTI_CACHE[jti] = AuthorizationCodeCacheEntry(exp_ts=exp_ts)


def consume_authorization_code_jti(jti: str) -> bool:
    _cleanup_cache()
    entry = _CODE_JTI_CACHE.get(jti)
    if not entry:
        return False
    _CODE_JTI_CACHE.pop(jti, None)
    return True


def normalize_next(next_url: str | None, *, base_url: str) -> str:

    if not next_url:
        return DEFAULT_NEXT_PATH
    if next_url.startswith('/') and not next_url.startswith('//'):
        return next_url
    try:
        target = urlparse(next_url)
        base = urlparse(base_url)
    except Exception:
        return DEFAULT_NEXT_PATH

    if target.scheme and target.netloc and (target.scheme, target.netloc) == (base.scheme, base.netloc):
        return next_url

    return DEFAULT_NEXT_PATH


def read_oauth_login_asset(relative_path: str) -> str:
    return (
        resources.files('amsdal_ml.mcp_server.mcp_oauth').joinpath(relative_path).read_text(encoding=TEXT_ENCODING_UTF8)
    )


def render_login_html(*, next_url: str, error: str | None = None) -> str:
    html = read_oauth_login_asset(LOGIN_TEMPLATE_PATH)
    error_block = f'<p class="error">{escape(error)}</p>' if error else ''
    return html.replace('{{next}}', escape(next_url, quote=True)).replace('{{error_block}}', error_block)


def pkce_s256(code_verifier: str) -> str:
    digest = hashlib.sha256(code_verifier.encode('utf-8')).digest()
    return base64.urlsafe_b64encode(digest).rstrip(b'=').decode('utf-8')


def get_signing_key() -> str:
    from amsdal.contrib.auth.errors import AuthenticationError
    from amsdal.contrib.auth.settings import auth_settings

    signing_key = auth_settings.AUTH_JWT_KEY
    if not signing_key:
        msg = 'JWT key is not set'
        raise AuthenticationError(msg)
    return signing_key


def decode_jwt_payload(token: str, signing_key: str) -> JwtDecodedPayload:
    decoded_value = jwt.decode(
        token,
        signing_key,
        algorithms=[JWT_ALGORITHM],
        options=Options(verify_aud=False),
    )
    return JwtDecodedPayload.model_validate(decoded_value)


def build_access_payload(
    *,
    sub: str,
    resource: str | None,
    scope: str | None,
    now: datetime,
) -> dict[str, str | int | None]:
    iat = int(now.timestamp())
    exp = int((now + timedelta(hours=ml_config.oauth_access_token_exp_hours)).timestamp())
    return {
        JwtClaims.ISSUER.value: ml_config.oauth_issuer,
        JwtClaims.SUBJECT.value: sub,
        JwtClaims.AUDIENCE.value: resource,
        JwtClaims.SCOPE.value: scope,
        JwtClaims.ISSUED_AT.value: iat,
        JwtClaims.EXPIRATION.value: exp,
        JwtClaims.JWT_ID.value: str(uuid.uuid4()),
        JwtClaims.TOKEN_TYPE.value: OAuthConstants.TOKEN_TYPE_ACCESS,
    }


def build_refresh_payload(
    *,
    sub: str,
    client_id: str,
    resource: str | None,
    scope: str | None,
    now: datetime,
) -> dict[str, str | int | None]:
    iat = int(now.timestamp())
    exp = int((now + timedelta(days=ml_config.oauth_refresh_token_exp_days)).timestamp())
    return {
        JwtClaims.ISSUER.value: ml_config.oauth_issuer,
        JwtClaims.SUBJECT.value: sub,
        JwtClaims.CLIENT_ID.value: client_id,
        JwtClaims.AUDIENCE.value: resource,
        JwtClaims.RESOURCE.value: resource,
        JwtClaims.SCOPE.value: scope,
        JwtClaims.ISSUED_AT.value: iat,
        JwtClaims.EXPIRATION.value: exp,
        JwtClaims.JWT_ID.value: str(uuid.uuid4()),
        JwtClaims.TOKEN_TYPE.value: OAuthConstants.TOKEN_TYPE_REFRESH,
    }


def access_token_expires_in_seconds() -> int:
    return int(timedelta(hours=ml_config.oauth_access_token_exp_hours).total_seconds())


def normalize_resource(resource: str | None) -> str | None:
    if resource is None:
        return None
    normalized = resource.strip()
    if not normalized:
        return None
    return normalized.rstrip('/')


def resources_match(left: str | None, right: str | None) -> bool:
    return normalize_resource(left) == normalize_resource(right)
