from datetime import UTC, datetime, timedelta
from uuid import uuid4

import jwt
from pydantic import BaseModel

from bluefox_auth.settings import AuthSettings


class TokenPayload(BaseModel):
    sub: int
    type: str  # "access", "refresh", "password_reset", "email_verify"
    exp: datetime
    iat: datetime
    jti: str
    audience: str


class TokenPair(BaseModel):
    access_token: str
    refresh_token: str
    token_type: str = "bearer"


class InvalidTokenError(Exception):
    pass


def create_access_token(user_id: int, settings: AuthSettings) -> str:
    now = datetime.now(UTC)
    payload = {
        "sub": str(user_id),
        "type": "access",
        "jti": uuid4().hex,
        "iat": now,
        "aud": settings.token_audience,
        "exp": now + timedelta(minutes=settings.access_token_expire_minutes),
    }
    return jwt.encode(payload, settings.secret_key, algorithm=settings.algorithm)


def create_refresh_token(user_id: int, settings: AuthSettings) -> tuple[str, str, str]:
    """Create a refresh token. Returns (token_string, jti, family_id)."""
    now = datetime.now(UTC)
    jti = uuid4().hex
    family_id = uuid4().hex
    payload = {
        "sub": str(user_id),
        "type": "refresh",
        "jti": jti,
        "iat": now,
        "aud": settings.token_audience,
        "exp": now + timedelta(days=settings.refresh_token_expire_days),
    }
    token = jwt.encode(payload, settings.secret_key, algorithm=settings.algorithm)
    return token, jti, family_id


def create_rotated_refresh_token(
    user_id: int, family_id: str, settings: AuthSettings
) -> tuple[str, str]:
    """Create a rotated refresh token reusing the family_id. Returns (token_string, jti)."""
    now = datetime.now(UTC)
    jti = uuid4().hex
    payload = {
        "sub": str(user_id),
        "type": "refresh",
        "jti": jti,
        "iat": now,
        "aud": settings.token_audience,
        "exp": now + timedelta(days=settings.refresh_token_expire_days),
    }
    token = jwt.encode(payload, settings.secret_key, algorithm=settings.algorithm)
    return token, jti


def create_token_pair(user_id: int, settings: AuthSettings) -> tuple[TokenPair, str, str]:
    """Create an access + refresh token pair.

    Returns (token_pair, refresh_jti, refresh_family_id).
    """
    access_token = create_access_token(user_id, settings)
    refresh_token, jti, family_id = create_refresh_token(user_id, settings)
    pair = TokenPair(access_token=access_token, refresh_token=refresh_token)
    return pair, jti, family_id


def decode_token(token: str, settings: AuthSettings, expected_type: str) -> TokenPayload:
    """Decode and validate a JWT token.

    expected_type is required — every decode call must declare what it expects.
    Raises InvalidTokenError on any failure.
    """
    try:
        payload = jwt.decode(
            token,
            settings.secret_key,
            algorithms=[settings.algorithm],
            audience=settings.token_audience,
        )
    except jwt.ExpiredSignatureError as err:
        raise InvalidTokenError("Token has expired") from err
    except jwt.InvalidAudienceError as err:
        raise InvalidTokenError("Invalid token audience") from err
    except jwt.PyJWTError as err:
        raise InvalidTokenError("Invalid token") from err

    if payload.get("type") != expected_type:
        raise InvalidTokenError(
            f"Expected token type '{expected_type}', got '{payload.get('type')}'"
        )

    try:
        return TokenPayload(
            sub=int(payload["sub"]),
            type=payload["type"],
            exp=datetime.fromtimestamp(payload["exp"], tz=UTC),
            iat=datetime.fromtimestamp(payload["iat"], tz=UTC),
            jti=payload["jti"],
            audience=payload["aud"],
        )
    except (KeyError, ValueError) as err:
        raise InvalidTokenError("Token is missing required claims") from err


def create_special_token(
    user_id: int,
    token_type: str,
    expire_minutes: int,
    settings: AuthSettings,
    extra_claims: dict | None = None,
) -> str:
    """Create a special-purpose token (password reset, email verification)."""
    now = datetime.now(UTC)
    payload = {
        "sub": str(user_id),
        "type": token_type,
        "jti": uuid4().hex,
        "iat": now,
        "aud": settings.token_audience,
        "exp": now + timedelta(minutes=expire_minutes),
    }
    if extra_claims:
        reserved = {"sub", "type", "jti", "iat", "aud", "exp"}
        overlap = reserved & extra_claims.keys()
        if overlap:
            raise ValueError(f"extra_claims may not override reserved fields: {overlap}")
        payload.update(extra_claims)
    return jwt.encode(payload, settings.secret_key, algorithm=settings.algorithm)
