use sea_orm::DatabaseConnection;
use time::OffsetDateTime;
use uuid::Uuid;

use crate::claims::{AccessClaims, RefreshClaims};
use crate::config::JwtSettings;
use crate::error::JwtError;
use crate::keys::SigningKeyCache;
use crate::session::{self, NewSession};
use crate::signer::{
    decode_access_token, decode_refresh_token, encode_access_token, encode_refresh_token,
};

pub struct TokenPair {
    pub access_token: String,
    pub refresh_token: String,
    pub access_claims: AccessClaims,
    pub refresh_claims: RefreshClaims,
    pub session_id: Uuid,
}

pub struct RotationContext<'a> {
    pub db: &'a DatabaseConnection,
    pub settings: &'a JwtSettings,
    pub cache: &'a SigningKeyCache,
}

/// Live user facts needed to mint a rotation pair with correct identity
/// claims. The auth layer reads these from `users`; digging them out of the
/// refresh token would let stale claims (admin/restricted/username) survive
/// rotation forever.
pub struct UserSnapshot {
    pub pwd_version: i64,
    pub username: String,
    pub is_admin: bool,
    pub is_restricted: bool,
}

fn current_key<'a>(
    ctx: &'a RotationContext,
) -> Result<(Uuid, &'a jsonwebtoken::EncodingKey, jsonwebtoken::Algorithm), JwtError> {
    let (kid, encoding) = ctx.cache.current_encoding().ok_or(JwtError::NoActiveKey)?;
    let alg = ctx
        .cache
        .get(&kid)
        .map(|k| k.algorithm)
        .ok_or(JwtError::NoActiveKey)?;
    Ok((kid, encoding, alg))
}

/// Issue a brand-new access/refresh pair for a fresh login (no parent).
/// `pwd_version` MUST equal the live `users.passwd_version` for the user
/// at the moment of issuance; the auth layer is responsible for reading
/// it. The value is embedded in both claims so any future refresh or
/// access-token check can reject the token if it lags behind a bump.
pub async fn issue_login_pair(
    ctx: &RotationContext<'_>,
    user_id: Uuid,
    username: String,
    is_admin: bool,
    is_restricted: bool,
    pwd_version: i64,
    app_id: Option<String>,
) -> Result<TokenPair, JwtError> {
    let (kid, encoding, alg) = current_key(ctx)?;
    let family_id = Uuid::new_v4();

    let access = AccessClaims::new(
        user_id,
        username.clone(),
        is_admin,
        is_restricted,
        pwd_version,
        ctx.settings.access_ttl.whole_seconds(),
        &ctx.settings.issuer,
        &ctx.settings.audience,
        &kid.to_string(),
    );
    let refresh = RefreshClaims::new(
        user_id,
        family_id,
        None,
        pwd_version,
        ctx.settings.refresh_ttl.whole_seconds(),
        &ctx.settings.issuer,
        &ctx.settings.audience,
        &kid.to_string(),
    );

    let access_token = encode_access_token(&access, encoding, alg)?;
    let refresh_token = encode_refresh_token(&refresh, encoding, alg)?;

    let session = session::create_session(
        ctx.db,
        NewSession {
            user_id,
            token_hash: session::hash_token(&refresh_token),
            family_id,
            parent_refresh_id: None,
            expires_at: OffsetDateTime::now_utc() + ctx.settings.refresh_ttl,
            grace_expires_at: None,
            signing_key_id: Some(kid),
            app_id,
        },
    )
    .await?;

    Ok(TokenPair {
        access_token,
        refresh_token,
        access_claims: access,
        refresh_claims: refresh,
        session_id: session.id,
    })
}

/// Rotate a refresh token: validate the presented one, issue a new pair, and
/// mark the old session replaced (with grace window) so replays are rejected.
/// `user` MUST be the live `users` row; tokens whose `pwd_version` claim lags
/// behind it are rejected here so a stolen refresh token cannot outlive a
/// password reset, and the new access claims are minted from `user` (not from
/// the presented token) so admin/restricted/username state cannot go stale.
pub async fn rotate_refresh(
    ctx: &RotationContext<'_>,
    refresh_token: &str,
    user: &UserSnapshot,
    app_id: Option<String>,
) -> Result<TokenPair, JwtError> {
    let kid = crate::signer::extract_kid(refresh_token).ok_or(JwtError::InvalidHeader)?;
    let cached = ctx
        .cache
        .get(&kid)
        .ok_or(JwtError::UnknownKid(kid.to_string()))?;

    let old_refresh = decode_refresh_token(
        refresh_token,
        &cached.decoding,
        cached.algorithm,
        &ctx.settings.issuer,
        &ctx.settings.audience,
        ctx.settings.clock_skew_seconds,
    )?;

    if old_refresh.pwd_version < user.pwd_version {
        return Err(JwtError::SessionRevoked);
    }

    let old_session = session::find_session_by_hash(ctx.db, &session::hash_token(refresh_token))
        .await?
        .ok_or(JwtError::SessionNotFound)?;

    if old_session.is_revoked {
        return Err(JwtError::SessionRevoked);
    }
    if old_session.replaced_by.is_some() {
        // Allow replays within the grace window.
        let in_grace = old_session
            .grace_expires_at
            .map(|t| t > OffsetDateTime::now_utc())
            .unwrap_or(false);
        if !in_grace {
            return Err(JwtError::SessionReplaced);
        }
    }

    let (new_kid, encoding, alg) = current_key(ctx)?;
    let new_refresh = RefreshClaims::new(
        old_refresh.sub,
        old_refresh.family_id,
        Some(old_refresh.jti),
        user.pwd_version,
        ctx.settings.refresh_ttl.whole_seconds(),
        &ctx.settings.issuer,
        &ctx.settings.audience,
        &new_kid.to_string(),
    );
    let new_access = AccessClaims::new(
        old_refresh.sub,
        user.username.clone(),
        user.is_admin,
        user.is_restricted,
        user.pwd_version,
        ctx.settings.access_ttl.whole_seconds(),
        &ctx.settings.issuer,
        &ctx.settings.audience,
        &new_kid.to_string(),
    );

    let new_refresh_token = encode_refresh_token(&new_refresh, encoding, alg)?;
    let new_access_token = encode_access_token(&new_access, encoding, alg)?;

    let new_session = session::create_session(
        ctx.db,
        NewSession {
            user_id: old_refresh.sub,
            token_hash: session::hash_token(&new_refresh_token),
            family_id: old_refresh.family_id,
            parent_refresh_id: Some(old_refresh.jti),
            expires_at: OffsetDateTime::now_utc() + ctx.settings.refresh_ttl,
            grace_expires_at: Some(OffsetDateTime::now_utc() + ctx.settings.grace_window),
            signing_key_id: Some(new_kid),
            app_id,
        },
    )
    .await?;

    session::mark_replaced(
        ctx.db,
        old_session.id,
        new_session.id,
        OffsetDateTime::now_utc() + ctx.settings.grace_window,
    )
    .await?;

    Ok(TokenPair {
        access_token: new_access_token,
        refresh_token: new_refresh_token,
        access_claims: new_access,
        refresh_claims: new_refresh,
        session_id: new_session.id,
    })
}

/// Validate an access token against the active key set.
pub async fn validate_access(
    ctx: &RotationContext<'_>,
    access_token: &str,
) -> Result<AccessClaims, JwtError> {
    let kid = crate::signer::extract_kid(access_token).ok_or(JwtError::InvalidHeader)?;
    let cached = ctx
        .cache
        .get(&kid)
        .ok_or(JwtError::UnknownKid(kid.to_string()))?;
    decode_access_token(
        access_token,
        &cached.decoding,
        cached.algorithm,
        &ctx.settings.issuer,
        &ctx.settings.audience,
        ctx.settings.clock_skew_seconds,
    )
}