use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::RwLock;

#[derive(Clone)]
struct Session {
    created_at: Instant,
}

#[derive(Clone)]
pub struct SessionStore {
    sessions: Arc<RwLock<HashMap<[u8; 32], Session>>>,
}

impl SessionStore {
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn create(&self) -> String {
        let token_bytes: [u8; 32] = rand::random();
        let token = hex::encode(token_bytes);
        let hash = hash_token(&token);
        let session = Session {
            created_at: Instant::now(),
        };
        self.sessions.write().await.insert(hash, session);
        token
    }

    pub async fn validate(&self, token: &str) -> bool {
        let hash = hash_token(token);
        let sessions = self.sessions.read().await;
        if let Some(session) = sessions.get(&hash) {
            session.created_at.elapsed().as_secs() < 86400
        } else {
            false
        }
    }

    pub async fn remove(&self, token: &str) {
        let hash = hash_token(token);
        self.sessions.write().await.remove(&hash);
    }
}

fn hash_token(token: &str) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(token.as_bytes());
    hasher.finalize().into()
}

/// Extract the session token from a Cookie header value.
pub fn extract_session_cookie(headers: &hyper::HeaderMap) -> Option<String> {
    headers
        .get("cookie")
        .and_then(|v| v.to_str().ok())
        .and_then(|cookies| {
            cookies.split(';').find_map(|c| {
                let c = c.trim();
                c.strip_prefix("session=").map(|v| v.to_string())
            })
        })
        .filter(|v| !v.is_empty())
}

/// Rate limiter for login attempts: max 5 failures per 60 seconds per IP.
#[derive(Clone)]
pub struct LoginRateLimiter {
    attempts: Arc<RwLock<HashMap<IpAddr, Vec<Instant>>>>,
}

const MAX_ATTEMPTS: usize = 5;
const WINDOW_SECS: u64 = 60;

impl LoginRateLimiter {
    pub fn new() -> Self {
        Self {
            attempts: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Check if a login attempt is allowed for this IP. Returns false if rate limited.
    pub async fn check(&self, ip: IpAddr) -> bool {
        let mut attempts = self.attempts.write().await;
        let now = Instant::now();
        let entry = attempts.entry(ip).or_insert_with(Vec::new);

        // Remove attempts older than the window
        entry.retain(|t| now.duration_since(*t).as_secs() < WINDOW_SECS);

        entry.len() < MAX_ATTEMPTS
    }

    /// Record a failed login attempt.
    pub async fn record_failure(&self, ip: IpAddr) {
        let mut attempts = self.attempts.write().await;
        let entry = attempts.entry(ip).or_insert_with(Vec::new);
        entry.push(Instant::now());
    }
}