//! Authentication store implementations.

use async_trait::async_trait;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set};
use std::sync::Arc;
use std::time::SystemTime;
use tideway::auth::{
    JwtIssuer, RefreshTokenStore, TokenIssuer, UserCreator, UserLoader, UserStore,
    PasswordResetStore, VerificationStore,
};
use tideway::{Result, TidewayError};
use uuid::Uuid;
use chrono::{Duration, Utc};

use crate::entities::{user, refresh_token_family, verification_token};

/// Database-backed user store implementation.
#[derive(Clone)]
pub struct DbUserStore {
    db: Arc<DatabaseConnection>,
    max_failed_attempts: i32,
    lockout_duration_minutes: i64,
}

impl DbUserStore {
    pub fn new(db: Arc<DatabaseConnection>) -> Self {
        Self {
            db,
            max_failed_attempts: 5,
            lockout_duration_minutes: 15,
        }
    }
}

#[async_trait]
impl UserStore for DbUserStore {
    type User = user::Model;

    async fn find_by_email(&self, email: &str) -> Result<Option<Self::User>> {
        let user = user::Entity::find()
            .filter(user::Column::Email.eq(email.to_lowercase()))
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(user)
    }

    async fn find_by_id(&self, id: &str) -> Result<Option<Self::User>> {
        let uuid = Uuid::parse_str(id)
            .map_err(|_| TidewayError::BadRequest("Invalid user ID".into()))?;
        let user = user::Entity::find_by_id(uuid)
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(user)
    }

    fn user_id(&self, user: &Self::User) -> String {
        user.id.to_string()
    }

    fn user_email(&self, user: &Self::User) -> String {
        user.email.clone()
    }

    fn user_name(&self, user: &Self::User) -> Option<String> {
        user.name.clone()
    }

    async fn get_password_hash(&self, user: &Self::User) -> Result<String> {
        Ok(user.password_hash.clone())
    }

    async fn update_password_hash(&self, user: &Self::User, hash: &str) -> Result<()> {
        let mut active: user::ActiveModel = user.clone().into();
        active.password_hash = Set(hash.to_string());
        active.updated_at = Set(Utc::now().into());
        active.update(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }

    async fn is_verified(&self, user: &Self::User) -> Result<bool> {
        Ok(user.email_verified_at.is_some())
    }

    async fn mark_verified(&self, user: &Self::User) -> Result<()> {
        let mut active: user::ActiveModel = user.clone().into();
        active.email_verified_at = Set(Some(Utc::now().into()));
        active.updated_at = Set(Utc::now().into());
        active.update(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }

    async fn is_locked(&self, user: &Self::User) -> Result<Option<SystemTime>> {
        if let Some(locked_until) = user.locked_until {
            let locked_time: SystemTime = locked_until.into();
            if locked_time > SystemTime::now() {
                return Ok(Some(locked_time));
            }
        }
        Ok(None)
    }

    async fn record_failed_attempt(&self, user: &Self::User) -> Result<()> {
        let new_count = user.failed_attempts + 1;
        let mut active: user::ActiveModel = user.clone().into();
        active.failed_attempts = Set(new_count);

        if new_count >= self.max_failed_attempts {
            let locked_until = Utc::now() + Duration::minutes(self.lockout_duration_minutes);
            active.locked_until = Set(Some(locked_until.into()));
        }

        active.updated_at = Set(Utc::now().into());
        active.update(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }

    async fn record_failed_mfa_attempt(&self, user: &Self::User) -> Result<()> {
        let new_count = user.failed_mfa_attempts + 1;
        let mut active: user::ActiveModel = user.clone().into();
        active.failed_mfa_attempts = Set(new_count);

        if new_count >= self.max_failed_attempts {
            let locked_until = Utc::now() + Duration::minutes(self.lockout_duration_minutes);
            active.locked_until = Set(Some(locked_until.into()));
        }

        active.updated_at = Set(Utc::now().into());
        active.update(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }

    async fn clear_failed_attempts(&self, user: &Self::User) -> Result<()> {
        let mut active: user::ActiveModel = user.clone().into();
        active.failed_attempts = Set(0);
        active.failed_mfa_attempts = Set(0);
        active.locked_until = Set(None);
        active.updated_at = Set(Utc::now().into());
        active.update(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }

    async fn has_mfa_enabled(&self, _user: &Self::User) -> Result<bool> {
        // TODO: Implement MFA check
        Ok(false)
    }

    async fn get_totp_secret(&self, _user: &Self::User) -> Result<Option<String>> {
        // TODO: Implement TOTP secret retrieval
        Ok(None)
    }

    async fn get_backup_codes(&self, _user: &Self::User) -> Result<Vec<String>> {
        // TODO: Implement backup codes retrieval
        Ok(vec![])
    }

    async fn remove_backup_code(&self, _user: &Self::User, _index: usize) -> Result<()> {
        // TODO: Implement backup code removal
        Ok(())
    }
}

#[async_trait]
impl UserCreator for DbUserStore {
    type User = user::Model;

    fn user_id(&self, user: &Self::User) -> String {
        user.id.to_string()
    }

    async fn email_exists(&self, email: &str) -> Result<bool> {
        let exists = user::Entity::find()
            .filter(user::Column::Email.eq(email.to_lowercase()))
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?
            .is_some();
        Ok(exists)
    }

    async fn create_user(
        &self,
        email: &str,
        password_hash: &str,
        name: Option<&str>,
    ) -> Result<Self::User> {
        let now = Utc::now();
        let auto_verify = std::env::var("AUTO_VERIFY_EMAIL").is_ok();

        let user = user::ActiveModel {
            id: Set(Uuid::new_v4()),
            email: Set(email.to_lowercase()),
            password_hash: Set(password_hash.to_string()),
            name: Set(name.map(String::from)),
            {{#if has_organizations}}
            organization_id: Set(None),
            {{/if}}
            email_verified_at: Set(if auto_verify { Some(now.into()) } else { None }),
            locked_until: Set(None),
            failed_attempts: Set(0),
            failed_mfa_attempts: Set(0),
            is_platform_admin: Set(false),
            created_at: Set(now.into()),
            updated_at: Set(now.into()),
        };

        let user = user.insert(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(user)
    }

    async fn send_verification_email(&self, _user: &Self::User) -> Result<()> {
        tracing::info!("Would send verification email to user");
        Ok(())
    }
}

#[async_trait]
impl PasswordResetStore for DbUserStore {
    type User = user::Model;

    async fn find_by_email(&self, email: &str) -> Result<Option<Self::User>> {
        UserStore::find_by_email(self, email).await
    }

    fn user_id(&self, user: &Self::User) -> String {
        UserStore::user_id(self, user)
    }

    async fn store_reset_token(
        &self,
        user_id: &str,
        token_hash: &str,
        expires: SystemTime,
    ) -> Result<()> {
        let uuid = Uuid::parse_str(user_id)
            .map_err(|_| TidewayError::BadRequest("Invalid user ID".into()))?;

        let expires_at: chrono::DateTime<Utc> = expires.into();

        let token = verification_token::ActiveModel {
            id: Set(Uuid::new_v4()),
            user_id: Set(uuid),
            token_hash: Set(token_hash.to_string()),
            token_type: Set("password_reset".to_string()),
            expires_at: Set(expires_at.into()),
            used_at: Set(None),
            created_at: Set(Utc::now().into()),
        };

        token.insert(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }

    async fn consume_reset_token(&self, token_hash: &str) -> Result<Option<String>> {
        let token = verification_token::Entity::find()
            .filter(verification_token::Column::TokenHash.eq(token_hash))
            .filter(verification_token::Column::TokenType.eq("password_reset"))
            .filter(verification_token::Column::UsedAt.is_null())
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;

        if let Some(token) = token {
            let expires_at: SystemTime = token.expires_at.into();
            if expires_at < SystemTime::now() {
                return Ok(None);
            }

            let user_id = token.user_id.to_string();
            let mut active: verification_token::ActiveModel = token.into();
            active.used_at = Set(Some(Utc::now().into()));
            active.update(self.db.as_ref()).await
                .map_err(|e| TidewayError::Database(e.to_string()))?;

            return Ok(Some(user_id));
        }
        Ok(None)
    }

    async fn update_password(&self, user_id: &str, hash: &str) -> Result<()> {
        let uuid = Uuid::parse_str(user_id)
            .map_err(|_| TidewayError::BadRequest("Invalid user ID".into()))?;

        let user = user::Entity::find_by_id(uuid)
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?
            .ok_or_else(|| TidewayError::NotFound("User not found".into()))?;

        self.update_password_hash(&user, hash).await
    }

    async fn invalidate_sessions(&self, _user_id: &str) -> Result<()> {
        Ok(())
    }

    async fn send_reset_email(
        &self,
        _user: &Self::User,
        _token: &str,
        _expires_in: std::time::Duration,
    ) -> Result<()> {
        tracing::info!("Would send password reset email");
        Ok(())
    }
}

#[async_trait]
impl VerificationStore for DbUserStore {
    async fn store_verification_token(
        &self,
        user_id: &str,
        token_hash: &str,
        expires: SystemTime,
    ) -> Result<()> {
        let uuid = Uuid::parse_str(user_id)
            .map_err(|_| TidewayError::BadRequest("Invalid user ID".into()))?;

        let expires_at: chrono::DateTime<Utc> = expires.into();

        let token = verification_token::ActiveModel {
            id: Set(Uuid::new_v4()),
            user_id: Set(uuid),
            token_hash: Set(token_hash.to_string()),
            token_type: Set("email_verification".to_string()),
            expires_at: Set(expires_at.into()),
            used_at: Set(None),
            created_at: Set(Utc::now().into()),
        };

        token.insert(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }

    async fn consume_verification_token(&self, token_hash: &str) -> Result<Option<String>> {
        let token = verification_token::Entity::find()
            .filter(verification_token::Column::TokenHash.eq(token_hash))
            .filter(verification_token::Column::TokenType.eq("email_verification"))
            .filter(verification_token::Column::UsedAt.is_null())
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;

        if let Some(token) = token {
            let expires_at: SystemTime = token.expires_at.into();
            if expires_at < SystemTime::now() {
                return Ok(None);
            }

            let user_id = token.user_id.to_string();
            let mut active: verification_token::ActiveModel = token.into();
            active.used_at = Set(Some(Utc::now().into()));
            active.update(self.db.as_ref()).await
                .map_err(|e| TidewayError::Database(e.to_string()))?;

            return Ok(Some(user_id));
        }
        Ok(None)
    }

    async fn mark_user_verified(&self, user_id: &str) -> Result<()> {
        let uuid = Uuid::parse_str(user_id)
            .map_err(|_| TidewayError::BadRequest("Invalid user ID".into()))?;

        let user = user::Entity::find_by_id(uuid)
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?
            .ok_or_else(|| TidewayError::NotFound("User not found".into()))?;

        self.mark_verified(&user).await
    }

    async fn send_verification_email(
        &self,
        _user_id: &str,
        _email: &str,
        _token: &str,
        _expires_in: std::time::Duration,
    ) -> Result<()> {
        tracing::info!("Would send verification email");
        Ok(())
    }
}

#[async_trait]
impl UserLoader for DbUserStore {
    type User = user::Model;

    async fn load_user(&self, user_id: &str) -> Result<Option<Self::User>> {
        self.find_by_id(user_id).await
    }

    fn user_email(&self, user: &Self::User) -> Option<String> {
        Some(user.email.clone())
    }

    fn user_name(&self, user: &Self::User) -> Option<String> {
        user.name.clone()
    }
}

/// Database-backed refresh token store implementation.
#[derive(Clone)]
pub struct DbRefreshTokenStore {
    db: Arc<DatabaseConnection>,
}

impl DbRefreshTokenStore {
    pub fn new(db: Arc<DatabaseConnection>) -> Self {
        Self { db }
    }
}

#[async_trait]
impl RefreshTokenStore for DbRefreshTokenStore {
    async fn is_family_revoked(&self, family: &str) -> Result<bool> {
        let record = refresh_token_family::Entity::find_by_id(family.to_string())
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;

        Ok(record.map(|r| r.revoked).unwrap_or(true))
    }

    async fn get_family_generation(&self, family: &str) -> Result<Option<u32>> {
        let record = refresh_token_family::Entity::find_by_id(family.to_string())
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;

        Ok(record.map(|r| r.generation as u32))
    }

    async fn set_family_generation(&self, family: &str, generation: u32) -> Result<()> {
        let record = refresh_token_family::Entity::find_by_id(family.to_string())
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;

        if let Some(record) = record {
            let mut active: refresh_token_family::ActiveModel = record.into();
            active.generation = Set(generation as i32);
            active.update(self.db.as_ref()).await
                .map_err(|e| TidewayError::Database(e.to_string()))?;
        }
        Ok(())
    }

    async fn revoke_family(&self, family: &str) -> Result<()> {
        let record = refresh_token_family::Entity::find_by_id(family.to_string())
            .one(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;

        if let Some(record) = record {
            let mut active: refresh_token_family::ActiveModel = record.into();
            active.revoked = Set(true);
            active.update(self.db.as_ref()).await
                .map_err(|e| TidewayError::Database(e.to_string()))?;
        }
        Ok(())
    }

    async fn revoke_all_for_user(&self, user_id: &str) -> Result<()> {
        let uuid = Uuid::parse_str(user_id)
            .map_err(|_| TidewayError::BadRequest("Invalid user ID".into()))?;

        let families = refresh_token_family::Entity::find()
            .filter(refresh_token_family::Column::UserId.eq(uuid))
            .filter(refresh_token_family::Column::Revoked.eq(false))
            .all(self.db.as_ref())
            .await
            .map_err(|e| TidewayError::Database(e.to_string()))?;

        for family in families {
            let mut active: refresh_token_family::ActiveModel = family.into();
            active.revoked = Set(true);
            active.update(self.db.as_ref()).await
                .map_err(|e| TidewayError::Database(e.to_string()))?;
        }
        Ok(())
    }

    async fn associate_family_with_user(&self, family: &str, user_id: &str) -> Result<()> {
        let uuid = Uuid::parse_str(user_id)
            .map_err(|_| TidewayError::BadRequest("Invalid user ID".into()))?;

        let record = refresh_token_family::ActiveModel {
            family: Set(family.to_string()),
            user_id: Set(uuid),
            generation: Set(0),
            revoked: Set(false),
            created_at: Set(Utc::now().into()),
        };

        record.insert(self.db.as_ref()).await
            .map_err(|e| TidewayError::Database(e.to_string()))?;
        Ok(())
    }
}

/// In-memory MFA token store (for development).
#[derive(Default)]
pub struct InMemoryMfaTokenStore {
    tokens: std::sync::RwLock<std::collections::HashMap<String, (String, std::time::Instant)>>,
}

impl InMemoryMfaTokenStore {
    pub fn new() -> Self {
        Self {
            tokens: std::sync::RwLock::new(std::collections::HashMap::new()),
        }
    }
}

#[async_trait]
impl tideway::auth::MfaTokenStore for InMemoryMfaTokenStore {
    async fn store(&self, token: &str, user_id: &str, ttl: std::time::Duration) -> Result<()> {
        let expires_at = std::time::Instant::now() + ttl;
        self.tokens.write().unwrap().insert(token.to_string(), (user_id.to_string(), expires_at));
        Ok(())
    }

    async fn consume(&self, token: &str) -> Result<Option<String>> {
        let mut tokens = self.tokens.write().unwrap();
        if let Some((user_id, expires_at)) = tokens.remove(token) {
            if std::time::Instant::now() < expires_at {
                return Ok(Some(user_id));
            }
        }
        Ok(None)
    }
}

/// Token issuer that adds custom claims to access tokens.
#[derive(Clone)]
pub struct AppTokenIssuer {
    jwt_issuer: Arc<JwtIssuer>,
}

impl AppTokenIssuer {
    pub fn new(jwt_issuer: Arc<JwtIssuer>) -> Self {
        Self { jwt_issuer }
    }
}

impl TokenIssuer for AppTokenIssuer {
    type User = user::Model;

    fn issue(&self, user: &Self::User, remember_me: bool) -> Result<tideway::auth::TokenIssuance> {
        use tideway::auth::TokenSubject;

        // Build custom claims
        #[derive(serde::Serialize)]
        struct CustomClaims {
            is_admin: bool,
            {{#if has_organizations}}
            #[serde(skip_serializing_if = "Option::is_none")]
            org_id: Option<String>,
            {{/if}}
        }

        let custom = CustomClaims {
            is_admin: user.is_platform_admin,
            {{#if has_organizations}}
            org_id: user.organization_id.clone(),
            {{/if}}
        };

        // Create token subject with user info
        let user_id = user.id.to_string();
        let mut subject = TokenSubject::new(&user_id)
            .with_email(&user.email)
            .with_custom(custom);

        // Add name if present
        if let Some(ref name) = user.name {
            subject = subject.with_name(name);
        }

        // Issue tokens
        let pair = self.jwt_issuer.issue(subject, remember_me)
            .map_err(|e| TidewayError::Internal(e.to_string()))?;

        Ok(tideway::auth::TokenIssuance {
            access_token: pair.access_token,
            refresh_token: pair.refresh_token,
            expires_in: pair.expires_in,
            family: pair.family,
        })
    }
}
