use chrono::{DateTime, Utc};
use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, info};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum Permission {
    CreateSession,
    ExecuteCommand,
    ViewSessions,
    DeleteSession,
    ManageAgents,
    SystemAdmin,
}

impl Permission {
    pub fn as_str(&self) -> &'static str {
        match self {
            Permission::CreateSession => "create_session",
            Permission::ExecuteCommand => "execute_command",
            Permission::ViewSessions => "view_sessions",
            Permission::DeleteSession => "delete_session",
            Permission::ManageAgents => "manage_agents",
            Permission::SystemAdmin => "system_admin",
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentCredentials {
    pub id: String,
    pub name: String,
    pub api_key: String,
    pub permissions: Vec<Permission>,
    pub created_at: DateTime<Utc>,
    pub last_used: Option<DateTime<Utc>>,
    pub active: bool,
    pub metadata: HashMap<String, String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct JwtClaims {
    pub sub: String, // agent_id
    pub name: String,
    pub permissions: Vec<String>,
    pub exp: usize,  // expiration time
    pub iat: usize,  // issued at
    pub jti: String, // JWT ID
}

pub struct AuthManager {
    agents: Arc<RwLock<HashMap<String, AgentCredentials>>>,
    jwt_secret: String,
    token_expiry: Duration,
}

impl AuthManager {
    pub fn new(jwt_secret: String) -> Self {
        Self {
            agents: Arc::new(RwLock::new(HashMap::new())),
            jwt_secret,
            token_expiry: Duration::from_secs(24 * 60 * 60), // 24 hours
        }
    }

    pub async fn add_agent(&self, agent: AgentCredentials) -> Result<(), anyhow::Error> {
        let mut agents = self.agents.write().await;
        if agents.contains_key(&agent.id) {
            return Err(anyhow::anyhow!("Agent with ID {} already exists", agent.id));
        }

        agents.insert(agent.id.clone(), agent.clone());
        info!("Added new agent: {}", agent.name);
        Ok(())
    }

    pub async fn remove_agent(&self, agent_id: &str) -> Result<(), anyhow::Error> {
        let mut agents = self.agents.write().await;
        if agents.remove(agent_id).is_none() {
            return Err(anyhow::anyhow!("Agent with ID {} not found", agent_id));
        }

        info!("Removed agent: {}", agent_id);
        Ok(())
    }

    pub async fn get_agent(&self, agent_id: &str) -> Option<AgentCredentials> {
        let agents = self.agents.read().await;
        agents.get(agent_id).cloned()
    }

    pub async fn list_agents(&self) -> Vec<AgentCredentials> {
        let agents = self.agents.read().await;
        agents.values().cloned().collect()
    }

    pub async fn validate_token(&self, token: &str) -> Result<String, anyhow::Error> {
        let token_data = decode::<JwtClaims>(
            token,
            &DecodingKey::from_secret(self.jwt_secret.as_ref()),
            &Validation::default(),
        )?;

        let claims = token_data.claims;
        let agent_id = claims.sub.clone();

        // Check if agent exists and is active
        let agent = self
            .get_agent(&agent_id)
            .await
            .ok_or_else(|| anyhow::anyhow!("Agent not found: {}", agent_id))?;

        if !agent.active {
            return Err(anyhow::anyhow!("Agent is inactive: {}", agent_id));
        }

        debug!("Successfully validated token for agent: {}", agent_id);
        Ok(agent_id)
    }

    pub async fn authenticate_api_key(&self, api_key: &str) -> Result<String, anyhow::Error> {
        let agents = self.agents.read().await;
        let agent = agents
            .values()
            .find(|a| a.api_key == api_key && a.active)
            .ok_or_else(|| anyhow::anyhow!("Invalid API key"))?;

        info!("Agent authenticated via API key: {}", agent.name);
        Ok(agent.id.clone())
    }

    pub async fn generate_token(&self, agent_id: &str) -> Result<String, anyhow::Error> {
        let agent = self
            .get_agent(agent_id)
            .await
            .ok_or_else(|| anyhow::anyhow!("Agent not found: {}", agent_id))?;

        let now = Utc::now();
        let exp = now
            + chrono::Duration::from_std(self.token_expiry)
                .map_err(|e| anyhow::anyhow!("Invalid expiration duration: {}", e))?;

        let claims = JwtClaims {
            sub: agent.id.clone(),
            name: agent.name.clone(),
            permissions: agent
                .permissions
                .iter()
                .map(|p| p.as_str().to_string())
                .collect(),
            exp: exp.timestamp() as usize,
            iat: now.timestamp() as usize,
            jti: Uuid::new_v4().to_string(),
        };

        let token = encode(
            &Header::default(),
            &claims,
            &EncodingKey::from_secret(self.jwt_secret.as_ref()),
        )?;

        info!("Generated JWT token for agent: {}", agent.name);
        Ok(token)
    }

    pub async fn check_permission(&self, agent_id: &str, permission: &Permission) -> bool {
        let agents = self.agents.read().await;
        if let Some(agent) = agents.get(agent_id) {
            agent.permissions.contains(permission)
                || agent.permissions.contains(&Permission::SystemAdmin)
        } else {
            false
        }
    }

    pub async fn update_last_used(&self, agent_id: &str) -> Result<(), anyhow::Error> {
        let mut agents = self.agents.write().await;
        if let Some(agent) = agents.get_mut(agent_id) {
            agent.last_used = Some(Utc::now());
        }
        Ok(())
    }

    pub async fn get_agent_stats(&self) -> AgentStats {
        let agents = self.agents.read().await;
        let total_agents = agents.len();
        let active_agents = agents.values().filter(|a| a.active).count();
        let agents_with_recent_activity = agents
            .values()
            .filter(|a| {
                if let Some(last_used) = a.last_used {
                    Utc::now() - last_used < chrono::Duration::hours(24)
                } else {
                    false
                }
            })
            .count();

        AgentStats {
            total_agents,
            active_agents,
            agents_with_recent_activity,
        }
    }

    // Initialize with a default admin agent for development
    pub async fn init_default_agent(&self) -> Result<String, anyhow::Error> {
        let admin_agent = AgentCredentials {
            id: "admin".to_string(),
            name: "Default Admin Agent".to_string(),
            api_key: "sk-ditto-admin-2024".to_string(),
            permissions: vec![
                Permission::CreateSession,
                Permission::ExecuteCommand,
                Permission::ViewSessions,
                Permission::DeleteSession,
                Permission::ManageAgents,
                Permission::SystemAdmin,
            ],
            created_at: Utc::now(),
            last_used: None,
            active: true,
            metadata: HashMap::new(),
        };

        let agent_id = admin_agent.id.clone();
        self.add_agent(admin_agent).await?;
        info!("Initialized default admin agent");

        // Generate a JWT token for the admin agent
        self.generate_token(&agent_id).await
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentStats {
    pub total_agents: usize,
    pub active_agents: usize,
    pub agents_with_recent_activity: usize,
}

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

    #[tokio::test]
    async fn test_agent_authentication() {
        let auth_manager = AuthManager::new("test_secret".to_string());

        // Add test agent
        let agent = AgentCredentials {
            id: "test_agent".to_string(),
            name: "Test Agent".to_string(),
            api_key: "test_key".to_string(),
            permissions: vec![Permission::CreateSession],
            created_at: Utc::now(),
            last_used: None,
            active: true,
            metadata: HashMap::new(),
        };

        auth_manager.add_agent(agent).await.unwrap();

        // Test API key authentication
        let result = auth_manager.authenticate_api_key("test_key").await;
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "test_agent");

        // Test invalid API key
        let result = auth_manager.authenticate_api_key("invalid_key").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_jwt_token_generation() {
        let auth_manager = AuthManager::new("test_secret".to_string());

        // Add test agent
        let agent = AgentCredentials {
            id: "test_agent".to_string(),
            name: "Test Agent".to_string(),
            api_key: "test_key".to_string(),
            permissions: vec![Permission::CreateSession],
            created_at: Utc::now(),
            last_used: None,
            active: true,
            metadata: HashMap::new(),
        };

        auth_manager.add_agent(agent).await.unwrap();

        // Generate token
        let token = auth_manager.generate_token("test_agent").await.unwrap();
        assert!(!token.is_empty());

        // Validate token
        let agent_id = auth_manager.validate_token(&token).await.unwrap();
        assert_eq!(agent_id, "test_agent");
    }

    #[tokio::test]
    async fn test_permission_checking() {
        let auth_manager = AuthManager::new("test_secret".to_string());

        let agent = AgentCredentials {
            id: "test_agent".to_string(),
            name: "Test Agent".to_string(),
            api_key: "test_key".to_string(),
            permissions: vec![Permission::CreateSession],
            created_at: Utc::now(),
            last_used: None,
            active: true,
            metadata: HashMap::new(),
        };

        auth_manager.add_agent(agent).await.unwrap();

        // Test valid permission
        assert!(
            auth_manager
                .check_permission("test_agent", &Permission::CreateSession)
                .await
        );

        // Test invalid permission
        assert!(
            !auth_manager
                .check_permission("test_agent", &Permission::DeleteSession)
                .await
        );

        // Test non-existent agent
        assert!(
            !auth_manager
                .check_permission("non_existent", &Permission::CreateSession)
                .await
        );
    }
}
