use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use jsonwebtoken::{DecodingKey, EncodingKey};
use sea_orm::{DatabaseConnection, EntityTrait, QueryFilter};
use time::OffsetDateTime;
use uuid::Uuid;

use crate::claims::sha256_hex;
use crate::error::JwtError;
use crate::signer::algorithm_from_str;

pub struct CachedKey {
    pub decoding: DecodingKey,
    pub algorithm: jsonwebtoken::Algorithm,
    pub is_rotating: bool,
    pub grace_ends_at: Option<OffsetDateTime>,
    pub expires_at: Option<OffsetDateTime>,
}

impl std::fmt::Debug for CachedKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CachedKey")
            .field("algorithm", &self.algorithm)
            .field("is_rotating", &self.is_rotating)
            .field("grace_ends_at", &self.grace_ends_at)
            .field("expires_at", &self.expires_at)
            .finish()
    }
}

pub struct SigningKeyCache {
    keys: HashMap<Uuid, CachedKey>,
    current_kid: Option<Uuid>,
    encoding: Option<EncodingKey>,
    pub last_loaded: Instant,
}

impl std::fmt::Debug for SigningKeyCache {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SigningKeyCache")
            .field("kids", &self.keys.keys().collect::<Vec<_>>())
            .field("current_kid", &self.current_kid)
            .field("has_encoding", &self.encoding.is_some())
            .field("last_loaded", &self.last_loaded)
            .finish()
    }
}

impl SigningKeyCache {
    pub fn new() -> Self {
        Self {
            keys: HashMap::new(),
            current_kid: None,
            encoding: None,
            last_loaded: Instant::now(),
        }
    }

    pub async fn load_from_db(&mut self, db: &DatabaseConnection) -> Result<(), JwtError> {
        use db::system::system_signing_keys::{Entity as Sk, SigningAlgorithm};
        let rows = Sk::find()
            .filter(Sk::COLUMN.is_active.eq(true))
            .filter(Sk::COLUMN.revoked_at.is_null())
            .all(db)
            .await?;

        self.keys.clear();
        self.current_kid = None;
        self.encoding = None;

        for row in rows {
            let alg = match row.algorithm {
                SigningAlgorithm::Ed25519 => jsonwebtoken::Algorithm::EdDSA,
                SigningAlgorithm::ES256 => jsonwebtoken::Algorithm::ES256,
                SigningAlgorithm::RS256 => jsonwebtoken::Algorithm::RS256,
            };
            let decoding = DecodingKey::from_ed_der(&row.public_key);
            let cached = CachedKey {
                decoding,
                algorithm: alg,
                is_rotating: row.is_rotating,
                grace_ends_at: row.grace_ends_at,
                expires_at: row.expires_at,
            };
            self.keys.insert(row.kid, cached);
            if row.is_signing && self.current_kid.is_none() {
                let private_key = row
                    .private_key_encrypted
                    .clone()
                    .ok_or(JwtError::NoActiveKey)?;
                let encoding = EncodingKey::from_ed_der(&private_key);
                self.encoding = Some(encoding);
                self.current_kid = Some(row.kid);
            }
        }

        self.last_loaded = Instant::now();
        Ok(())
    }

    pub fn get(&self, kid: &Uuid) -> Option<&CachedKey> {
        self.keys.get(kid)
    }

    pub fn current_kid(&self) -> Option<Uuid> {
        self.current_kid
    }

    pub fn current_encoding(&self) -> Option<(Uuid, &EncodingKey)> {
        let kid = self.current_kid?;
        let enc = self.encoding.as_ref()?;
        Some((kid, enc))
    }

    pub fn in_grace_window(&self, kid: &Uuid) -> bool {
        match self.keys.get(kid) {
            Some(k) => {
                k.is_rotating
                    && k.grace_ends_at
                        .map(|t| t > OffsetDateTime::now_utc())
                        .unwrap_or(false)
            }
            None => false,
        }
    }

    pub fn key_hash(&self, public_key: &[u8]) -> String {
        sha256_hex(public_key)
    }
}

impl Default for SigningKeyCache {
    fn default() -> Self {
        Self::new()
    }
}

pub type ArcSigningKeyCache = Arc<tokio::sync::RwLock<SigningKeyCache>>;

pub fn algorithm_from_str_name(s: &str) -> Result<jsonwebtoken::Algorithm, JwtError> {
    algorithm_from_str(s)
}