use jsonwebtoken::{
    decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation,
};
use serde::{Deserialize, Serialize};

use crate::error::{AuthError, Result};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum SubjectType {
    Human,
    Agent,
}

impl SubjectType {
    pub fn as_str(&self) -> &'static str {
        match self {
            SubjectType::Human => "human",
            SubjectType::Agent => "agent",
        }
    }
}

impl Default for SubjectType {
    fn default() -> Self {
        SubjectType::Human
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Claims {
    pub sub: String,
    pub exp: i64,
    pub iat: i64,
    #[serde(default)]
    pub session_id: Option<String>,
    #[serde(default)]
    pub subject_type: Option<SubjectType>,
}

pub fn encode_token(claims: &Claims, secret: &[u8]) -> Result<String> {
    encode(
        &Header::new(Algorithm::HS256),
        claims,
        &EncodingKey::from_secret(secret),
    )
    .map_err(|e| AuthError::TokenError(e.to_string()))
}

pub fn decode_token(token: &str, secret: &[u8]) -> Result<Claims> {
    let token_data = decode::<Claims>(
        token,
        &DecodingKey::from_secret(secret),
        &Validation::new(Algorithm::HS256),
    )
    .map_err(|_| AuthError::InvalidToken)?;

    Ok(token_data.claims)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_claims() -> Claims {
        Claims {
            sub: "user:abc".into(),
            exp: 9_999_999_999,
            iat: 1_700_000_000,
            session_id: Some("session:abc".into()),
            subject_type: Some(SubjectType::Human),
        }
    }

    #[test]
    fn encode_then_decode_round_trip() {
        let secret = b"test-secret";
        let claims = make_claims();
        let token = encode_token(&claims, secret).expect("encode");
        let decoded = decode_token(&token, secret).expect("decode");
        assert_eq!(decoded.sub, claims.sub);
        assert_eq!(decoded.session_id, claims.session_id);
        assert_eq!(decoded.subject_type, claims.subject_type);
    }

    #[test]
    fn decode_with_wrong_secret_fails() {
        let claims = make_claims();
        let token = encode_token(&claims, b"secret-a").unwrap();
        assert!(decode_token(&token, b"secret-b").is_err());
    }

    #[test]
    fn claims_deserialize_old_tokens_without_subject_type() {
        let old_claims = serde_json::json!({
            "sub": "user:legacy",
            "exp": 1,
            "iat": 1,
            "session_id": "session:legacy"
        });

        let claims: Claims =
            serde_json::from_value(old_claims).expect("old claims should deserialize");
        assert_eq!(claims.sub, "user:legacy");
        assert_eq!(claims.subject_type, None);
    }

    #[test]
    fn claims_deserialize_new_tokens_with_subject_type() {
        let new_claims = serde_json::json!({
            "sub": "user:new",
            "exp": 1,
            "iat": 1,
            "session_id": "session:new",
            "subject_type": "human"
        });

        let claims: Claims =
            serde_json::from_value(new_claims).expect("new claims should deserialize");
        assert_eq!(claims.subject_type, Some(SubjectType::Human));
    }
}
