from __future__ import annotations

import time
from typing import Any, cast

import jwt
import structlog

from .exceptions import TokenExpiredError, TokenInvalidError, TokenValidationError
from .models import PackageMapToken, PackageToken, TajToken, Token
from .package_map import parse_s3_path
from .quilt_uri import validate_quilt_uri

logger = structlog.get_logger(__name__)


def create_token(
    subject: str,
    scopes: list[str],
    ttl: int,
    secret: str,
    issuer: str | None = None,
    audience: str | list[str] | None = None,
) -> str:
    """Create a signed JWT containing scope claims."""
    issued_at = int(time.time())
    expires_at = issued_at + ttl
    payload = {
        "sub": subject,
        "scopes": scopes,
        "iat": issued_at,
        "exp": expires_at,
    }
    if issuer:
        payload["iss"] = issuer
    if audience:
        payload["aud"] = audience
    return jwt.encode(payload, secret, algorithm="HS256")


def create_token_with_grants(
    subject: str,
    grants: list[str],
    ttl: int,
    secret: str,
    issuer: str | None = None,
    audience: str | list[str] | None = None,
) -> str:
    """Create a signed JWT containing grant claims."""
    issued_at = int(time.time())
    expires_at = issued_at + ttl
    payload = {
        "sub": subject,
        "grants": grants,
        "iat": issued_at,
        "exp": expires_at,
    }
    if issuer:
        payload["iss"] = issuer
    if audience:
        payload["aud"] = audience
    return jwt.encode(payload, secret, algorithm="HS256")


def create_token_with_package_grant(
    subject: str,
    quilt_uri: str,
    mode: str,
    ttl: int,
    secret: str,
    issuer: str | None = None,
    audience: str | list[str] | None = None,
) -> str:
    """Create a signed JWT containing a package grant."""
    if mode != "read":
        raise ValueError("package tokens only support read mode")
    issued_at = int(time.time())
    expires_at = issued_at + ttl
    payload = {
        "sub": subject,
        "quilt_uri": quilt_uri,
        "mode": mode,
        "iat": issued_at,
        "exp": expires_at,
    }
    if issuer:
        payload["iss"] = issuer
    if audience:
        payload["aud"] = audience
    return jwt.encode(payload, secret, algorithm="HS256")


def create_token_with_package_map(
    subject: str,
    quilt_uri: str,
    mode: str,
    ttl: int,
    secret: str,
    logical_bucket: str | None = None,
    logical_key: str | None = None,
    logical_s3_path: str | None = None,
    issuer: str | None = None,
    audience: str | list[str] | None = None,
) -> str:
    """Create a signed JWT containing a package map translation grant."""
    if mode != "read":
        raise ValueError("package tokens only support read mode")
    issued_at = int(time.time())
    expires_at = issued_at + ttl
    payload = {
        "sub": subject,
        "quilt_uri": quilt_uri,
        "mode": mode,
        "iat": issued_at,
        "exp": expires_at,
    }
    if logical_s3_path:
        payload["logical_s3_path"] = logical_s3_path
    if logical_bucket:
        payload["logical_bucket"] = logical_bucket
    if logical_key:
        payload["logical_key"] = logical_key
    if issuer:
        payload["iss"] = issuer
    if audience:
        payload["aud"] = audience
    return jwt.encode(payload, secret, algorithm="HS256")


def create_taj_token(
    subject: str,
    grants: list[str],
    manifest_hash: str,
    package_name: str,
    registry: str,
    ttl: int,
    secret: str,
    issuer: str | None = None,
    audience: str | list[str] | None = None,
) -> str:
    """Create a signed JWT containing RALE TAJ claims."""
    issued_at = int(time.time())
    expires_at = issued_at + ttl
    payload = {
        "sub": subject,
        "grants": grants,
        "manifest_hash": manifest_hash,
        "package_name": package_name,
        "registry": registry,
        "iat": issued_at,
        "exp": expires_at,
    }
    if issuer:
        payload["iss"] = issuer
    if audience:
        payload["aud"] = audience
    return jwt.encode(payload, secret, algorithm="HS256")


def validate_package_token(token_str: str, secret: str) -> PackageToken:
    """Validate a JWT signature and return a decoded PackageToken model."""
    try:
        payload = jwt.decode(token_str, secret, algorithms=["HS256"])
    except jwt.ExpiredSignatureError as exc:
        logger.warning("package_token_expired", error=str(exc))
        raise TokenExpiredError("token expired") from exc
    except jwt.InvalidTokenError as exc:
        logger.warning("package_token_invalid", error=str(exc))
        raise TokenInvalidError("invalid token") from exc
    except Exception as exc:
        logger.error("unexpected_package_token_validation_error", error=str(exc), exc_info=True)
        raise TokenValidationError(f"unexpected token validation error: {exc}") from exc

    subject = payload.get("sub")
    if not isinstance(subject, str) or not subject.strip():
        raise TokenValidationError("token subject is required")

    quilt_uri = payload.get("quilt_uri")
    if not isinstance(quilt_uri, str) or not quilt_uri.strip():
        raise TokenValidationError("token quilt_uri is required")

    try:
        quilt_uri = validate_quilt_uri(quilt_uri)
    except ValueError as exc:
        raise TokenValidationError(f"invalid quilt uri: {exc}") from exc

    mode = payload.get("mode")
    if mode != "read":
        raise TokenValidationError("token mode must be 'read'")

    try:
        return PackageToken(
            subject=subject,
            quilt_uri=quilt_uri,
            mode=mode,
            issued_at=int(payload.get("iat", 0)),
            expires_at=int(payload.get("exp", 0)),
        )
    except Exception as exc:
        logger.error("package_token_model_creation_failed", error=str(exc), exc_info=True)
        raise TokenValidationError(f"failed to create token model: {exc}") from exc


def validate_package_map_token(token_str: str, secret: str) -> PackageMapToken:
    """Validate a JWT signature and return a decoded PackageMapToken model."""
    try:
        payload = jwt.decode(token_str, secret, algorithms=["HS256"])
    except jwt.ExpiredSignatureError as exc:
        logger.warning("package_map_token_expired", error=str(exc))
        raise TokenExpiredError("token expired") from exc
    except jwt.InvalidTokenError as exc:
        logger.warning("package_map_token_invalid", error=str(exc))
        raise TokenInvalidError("invalid token") from exc
    except Exception as exc:
        logger.error("unexpected_package_map_token_validation_error", error=str(exc), exc_info=True)
        raise TokenValidationError(f"unexpected token validation error: {exc}") from exc

    subject = payload.get("sub")
    if not isinstance(subject, str) or not subject.strip():
        raise TokenValidationError("token subject is required")

    quilt_uri = payload.get("quilt_uri")
    if not isinstance(quilt_uri, str) or not quilt_uri.strip():
        raise TokenValidationError("token quilt_uri is required")

    try:
        quilt_uri = validate_quilt_uri(quilt_uri)
    except ValueError as exc:
        raise TokenValidationError(f"invalid quilt uri: {exc}") from exc

    mode = payload.get("mode")
    if mode != "read":
        raise TokenValidationError("token mode must be 'read'")

    logical_bucket = payload.get("logical_bucket")
    logical_key = payload.get("logical_key")
    logical_s3_path = payload.get("logical_s3_path")
    if logical_s3_path:
        try:
            parsed_bucket, parsed_key = parse_s3_path(str(logical_s3_path))
        except ValueError as exc:
            raise TokenValidationError(f"invalid logical_s3_path: {exc}") from exc
        if logical_bucket and logical_bucket != parsed_bucket:
            raise TokenValidationError("logical_bucket does not match logical_s3_path")
        if logical_key and logical_key != parsed_key:
            raise TokenValidationError("logical_key does not match logical_s3_path")
        logical_bucket = parsed_bucket
        logical_key = parsed_key

    if not isinstance(logical_bucket, str) or not logical_bucket.strip():
        raise TokenValidationError("token logical_bucket is required")
    if not isinstance(logical_key, str) or not logical_key.strip():
        raise TokenValidationError("token logical_key is required")

    try:
        return PackageMapToken(
            subject=subject,
            quilt_uri=quilt_uri,
            mode=mode,
            logical_bucket=logical_bucket,
            logical_key=logical_key,
            issued_at=int(payload.get("iat", 0)),
            expires_at=int(payload.get("exp", 0)),
        )
    except Exception as exc:
        logger.error("package_map_token_model_creation_failed", error=str(exc), exc_info=True)
        raise TokenValidationError(f"failed to create token model: {exc}") from exc


def validate_taj_token(token_str: str, secret: str) -> TajToken:
    """Validate a JWT signature and return a decoded TajToken model."""
    try:
        payload = jwt.decode(token_str, secret, algorithms=["HS256"])
    except jwt.ExpiredSignatureError as exc:
        logger.warning("taj_token_expired", error=str(exc))
        raise TokenExpiredError("token expired") from exc
    except jwt.InvalidTokenError as exc:
        logger.warning("taj_token_invalid", error=str(exc))
        raise TokenInvalidError("invalid token") from exc
    except Exception as exc:
        logger.error("unexpected_taj_token_validation_error", error=str(exc), exc_info=True)
        raise TokenValidationError(f"unexpected token validation error: {exc}") from exc

    subject = payload.get("sub")
    if not isinstance(subject, str) or not subject.strip():
        raise TokenValidationError("token subject is required")

    grants = payload.get("grants")
    if grants is None:
        raise TokenValidationError("token grants are required")
    if not isinstance(grants, list):
        raise TokenValidationError("token grants must be a list")

    manifest_hash = payload.get("manifest_hash")
    if not isinstance(manifest_hash, str) or not manifest_hash.strip():
        raise TokenValidationError("token manifest_hash is required")

    package_name = payload.get("package_name")
    if not isinstance(package_name, str) or not package_name.strip():
        raise TokenValidationError("token package_name is required")

    registry = payload.get("registry")
    if not isinstance(registry, str) or not registry.strip():
        raise TokenValidationError("token registry is required")

    try:
        return TajToken(
            subject=subject,
            grants=grants,
            manifest_hash=manifest_hash,
            package_name=package_name,
            registry=registry,
            issued_at=int(payload.get("iat", 0)),
            expires_at=int(payload.get("exp", 0)),
        )
    except Exception as exc:
        logger.error("taj_token_model_creation_failed", error=str(exc), exc_info=True)
        raise TokenValidationError(f"failed to create token model: {exc}") from exc


def validate_token(token_str: str, secret: str) -> Token:
    """Validate a JWT signature and return the decoded Token model.

    Args:
        token_str: JWT token string to validate
        secret: Secret key used to verify the token signature

    Returns:
        Token model with decoded claims

    Raises:
        TokenExpiredError: If the token has expired
        TokenInvalidError: If the token is malformed or signature is invalid
        TokenValidationError: If token validation fails for other reasons
    """
    try:
        payload = jwt.decode(token_str, secret, algorithms=["HS256"])
    except jwt.ExpiredSignatureError as exc:
        logger.warning("token_expired", error=str(exc))
        raise TokenExpiredError("token expired") from exc
    except jwt.InvalidTokenError as exc:
        logger.warning("token_invalid", error=str(exc))
        raise TokenInvalidError("invalid token") from exc
    except Exception as exc:
        logger.error("unexpected_token_validation_error", error=str(exc), exc_info=True)
        raise TokenValidationError(f"unexpected token validation error: {exc}") from exc

    subject = payload.get("sub")
    if not isinstance(subject, str) or not subject.strip():
        raise TokenValidationError("token subject is required")

    scopes = payload.get("scopes")
    if scopes is None:
        raise TokenValidationError("token scopes are required")
    if not isinstance(scopes, list):
        raise TokenValidationError("token scopes must be a list")

    try:
        return Token(
            subject=subject,
            scopes=scopes,
            issued_at=int(payload.get("iat", 0)),
            expires_at=int(payload.get("exp", 0)),
        )
    except Exception as exc:
        logger.error("token_model_creation_failed", error=str(exc), exc_info=True)
        raise TokenValidationError(f"failed to create token model: {exc}") from exc


def decode_token(token_str: str) -> dict[str, Any]:
    """Decode a JWT without validating signature or expiration.

    Args:
        token_str: JWT token string to decode

    Returns:
        Dictionary containing the decoded JWT payload

    Raises:
        TokenInvalidError: If the token cannot be decoded
    """
    try:
        payload = jwt.decode(
            token_str,
            options={"verify_signature": False, "verify_exp": False},
            algorithms=["HS256"],
        )
        return cast(dict[str, Any], payload)
    except jwt.InvalidTokenError as exc:
        logger.warning("token_decode_failed", error=str(exc))
        raise TokenInvalidError(f"failed to decode token: {exc}") from exc
    except Exception as exc:
        logger.error("unexpected_decode_error", error=str(exc), exc_info=True)
        raise TokenInvalidError(f"unexpected error decoding token: {exc}") from exc


def is_expired(token: Token) -> bool:
    """Return True if the token is expired relative to current time."""
    return int(time.time()) >= token.expires_at
