"""Utilities for JWT token and session storage operations."""

# Standard library imports
import os
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Optional

# Third-party imports
import jwt
import structlog

# Local imports
from .session_storage import session_storage
from .token_blacklist import token_blacklist

logger = structlog.get_logger(__name__)

# JWT configuration from environment variables
JWT_SECRET_KEY = os.getenv(
    "JWT_SECRET_KEY",
    "your_super_secret_jwt_key_minimum_32_characters_long",
)
JWT_ALGORITHM = os.getenv("JWT_ALGORITHM", "HS256")
JWT_ACCESS_TOKEN_EXPIRE_MINUTES = int(
    os.getenv("JWT_ACCESS_TOKEN_EXPIRE_MINUTES", "60")
)
JWT_REFRESH_TOKEN_EXPIRE_DAYS = int(
    os.getenv("JWT_REFRESH_TOKEN_EXPIRE_DAYS", "7")
)


def create_jwt_token(
    user_info: Dict[str, Any], session_id: str, token_type: str = "access"
) -> str:
    """Create JWT token with session_id instead of Guacamole token.

    Args:
        user_info: User information dictionary
        session_id: Session ID in Redis
        token_type: Token type ("access" or "refresh")

    Returns:
        JWT token as string

    Raises:
        Exception: If token creation fails
    """
    try:
        if token_type == "refresh":
            expire_delta = timedelta(days=JWT_REFRESH_TOKEN_EXPIRE_DAYS)
        else:
            expire_delta = timedelta(
                minutes=JWT_ACCESS_TOKEN_EXPIRE_MINUTES
            )

        now = datetime.now(timezone.utc)
        payload = {
            "username": user_info["username"],
            "role": user_info["role"],
            "permissions": user_info.get("permissions", []),
            "session_id": session_id,
            "token_type": token_type,
            "exp": now + expire_delta,
            "iat": now,
            "iss": "remote-access-api",
        }

        optional_fields = [
            "full_name",
            "email",
            "organization",
            "organizational_role",
        ]
        for field in optional_fields:
            if field in user_info:
                payload[field] = user_info[field]

        token = jwt.encode(payload, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)

        logger.info(
            "JWT token created successfully",
            username=user_info["username"],
            token_type=token_type,
            session_id=session_id,
            expires_in_minutes=expire_delta.total_seconds() / 60,
            payload_keys=list(payload.keys()),
            token_prefix=token[:30] + "...",
        )

        return token

    except Exception as e:
        logger.error(
            "Failed to create JWT token",
            username=user_info.get("username", "unknown"),
            error=str(e),
        )
        raise


def verify_jwt_token(token: str) -> Optional[Dict[str, Any]]:
    """Verify and decode JWT token with blacklist check.

    Args:
        token: JWT token to verify

    Returns:
        Decoded payload or None if token is invalid
    """
    try:
        logger.debug("Starting JWT verification", token_prefix=token[:30] + "...")

        if token_blacklist.is_token_revoked(token):
            logger.info("JWT token is revoked", token_prefix=token[:20] + "...")
            return None

        logger.debug("Token not in blacklist, attempting decode")

        payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])

        logger.info(
            "JWT decode successful",
            username=payload.get("username"),
            payload_keys=list(payload.keys()),
            has_session_id="session_id" in payload,
            session_id=payload.get("session_id", "NOT_FOUND"),
        )

        required_fields = ["username", "role", "session_id", "exp", "iat"]
        for field in required_fields:
            if field not in payload:
                logger.warning(
                    "JWT token missing required field",
                    field=field,
                    username=payload.get("username", "unknown"),
                    available_fields=list(payload.keys()),
                )
                return None

        logger.debug("All required fields present")

        exp_timestamp = payload["exp"]
        current_timestamp = datetime.now(timezone.utc).timestamp()
        if current_timestamp > exp_timestamp:
            logger.info(
                "JWT token expired",
                username=payload["username"],
                expired_at=datetime.fromtimestamp(
                    exp_timestamp, tz=timezone.utc
                ).isoformat(),
                current_time=datetime.now(timezone.utc).isoformat(),
            )
            return None

        logger.debug(
            "Token not expired, checking Redis session",
            session_id=payload["session_id"],
        )

        session_data = session_storage.get_session(payload["session_id"])
        if not session_data:
            logger.warning(
                "Session not found for JWT token",
                username=payload["username"],
                session_id=payload["session_id"],
                possible_reasons=[
                    "session expired in Redis",
                    "session never created",
                    "Redis connection issue",
                ],
            )
            return None

        logger.debug(
            "Session found in Redis",
            username=payload["username"],
            session_id=payload["session_id"],
            session_keys=list(session_data.keys()),
        )

        if "guac_token" not in session_data:
            logger.error(
                "Session exists but missing guac_token",
                username=payload["username"],
                session_id=payload["session_id"],
                session_keys=list(session_data.keys()),
            )
            return None

        payload["guac_token"] = session_data["guac_token"]

        logger.info(
            "JWT token verified successfully",
            username=payload["username"],
            role=payload["role"],
            token_type=payload.get("token_type", "access"),
            session_id=payload["session_id"],
            guac_token_length=len(session_data["guac_token"]),
        )

        return payload

    except jwt.ExpiredSignatureError:
        logger.info(
            "JWT token expired (ExpiredSignatureError)",
            token_prefix=token[:20] + "...",
        )
        return None
    except jwt.InvalidTokenError as e:
        logger.warning(
            "Invalid JWT token (InvalidTokenError)",
            error=str(e),
            error_type=type(e).__name__,
            token_prefix=token[:20] + "...",
        )
        return None
    except Exception as e:
        logger.error(
            "Unexpected error verifying JWT token",
            error=str(e),
            error_type=type(e).__name__,
        )
        return None


def create_refresh_token(
    user_info: Dict[str, Any], session_id: str
) -> str:
    """Create refresh token.

    Args:
        user_info: User information dictionary
        session_id: Session ID in Redis

    Returns:
        Refresh token
    """
    return create_jwt_token(user_info, session_id, token_type="refresh")


def extract_token_from_header(
    authorization_header: Optional[str],
) -> Optional[str]:
    """Extract token from Authorization header.

    Args:
        authorization_header: Authorization header value

    Returns:
        JWT token or None
    """
    if not authorization_header:
        return None

    if not authorization_header.startswith("Bearer "):
        return None

    return authorization_header.split(" ", 1)[1]


def get_token_expiry_info(token: str) -> Optional[Dict[str, Any]]:
    """Get token expiration information.

    Args:
        token: JWT token

    Returns:
        Expiration information or None
    """
    try:
        payload = jwt.decode(token, options={"verify_signature": False})

        exp_timestamp = payload.get("exp")
        iat_timestamp = payload.get("iat")

        if not exp_timestamp:
            return None

        exp_datetime = datetime.fromtimestamp(exp_timestamp, tz=timezone.utc)
        iat_datetime = (
            datetime.fromtimestamp(iat_timestamp, tz=timezone.utc)
            if iat_timestamp
            else None
        )
        current_time = datetime.now(timezone.utc)

        return {
            "expires_at": exp_datetime.isoformat(),
            "issued_at": iat_datetime.isoformat() if iat_datetime else None,
            "expires_in_seconds": max(
                0, int((exp_datetime - current_time).total_seconds())
            ),
            "is_expired": current_time > exp_datetime,
            "username": payload.get("username"),
            "token_type": payload.get("token_type", "access"),
            "session_id": payload.get("session_id"),
        }

    except Exception as e:
        logger.error("Failed to get token expiry info", error=str(e))
        return None


def is_token_expired(token: str) -> bool:
    """Check if token is expired.

    Args:
        token: JWT token

    Returns:
        True if token is expired, False if valid
    """
    expiry_info = get_token_expiry_info(token)
    return expiry_info["is_expired"] if expiry_info else True


def revoke_token(
    token: str,
    reason: str = "logout",
    revoked_by: Optional[str] = None,
) -> bool:
    """Revoke token (add to blacklist).

    Args:
        token: JWT token to revoke
        reason: Revocation reason
        revoked_by: Who revoked the token

    Returns:
        True if token successfully revoked
    """
    try:
        return token_blacklist.revoke_token(token, reason, revoked_by)
    except Exception as e:
        logger.error("Failed to revoke token", error=str(e))
        return False


def revoke_user_tokens(
    username: str,
    reason: str = "force_logout",
    revoked_by: Optional[str] = None,
) -> int:
    """Revoke all user tokens.

    Args:
        username: Username
        reason: Revocation reason
        revoked_by: Who revoked the tokens

    Returns:
        Number of revoked tokens
    """
    try:
        return token_blacklist.revoke_user_tokens(
            username, reason, revoked_by
        )
    except Exception as e:
        logger.error(
            "Failed to revoke user tokens", username=username, error=str(e)
        )
        return 0 