use anyhow::{anyhow, Result};
use secp256k1::{Message, PublicKey, Secp256k1, ecdsa::Signature};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

/// Nonce response structure
#[derive(Debug, Serialize, Deserialize)]
pub struct NonceResponse {
    pub nonce: String,
    pub timestamp: u64,
}

/// Authentication manager for handling nonces and signature verification
#[derive(Debug)]
pub struct AuthManager {
    nonces: HashMap<String, NonceResponse>,
    authenticated_keys: HashMap<String, u64>, // pubkey -> timestamp
}

impl AuthManager {
    /// Create a new authentication manager
    pub fn new() -> Self {
        Self {
            nonces: HashMap::new(),
            authenticated_keys: HashMap::new(),
        }
    }

    /// Generate a nonce for a given public key
    pub fn generate_nonce(&mut self, pubkey: &str) -> String {
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();
        
        let nonce = format!("{}:{}", pubkey, timestamp);
        let nonce_hash = format!("{:x}", Sha256::digest(nonce.as_bytes()));
        
        self.nonces.insert(
            pubkey.to_string(),
            NonceResponse {
                nonce: nonce_hash.clone(),
                timestamp,
            },
        );
        
        nonce_hash
    }

    /// Verify a signature against a stored nonce
    pub fn verify_signature(&mut self, pubkey: &str, signature: &str) -> Result<bool> {
        // Get the nonce for this public key
        let nonce_response = self
            .nonces
            .get(pubkey)
            .ok_or_else(|| anyhow!("No nonce found for public key"))?;

        // Check if nonce is not too old (5 minutes)
        let current_time = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();
        
        if current_time - nonce_response.timestamp > 300 {
            return Err(anyhow!("Nonce expired"));
        }

        // Parse the public key
        let pubkey_bytes = hex::decode(pubkey)
            .map_err(|_| anyhow!("Invalid public key format"))?;
        
        let secp = Secp256k1::new();
        let public_key = PublicKey::from_slice(&pubkey_bytes)
            .map_err(|_| anyhow!("Invalid public key"))?;

        // Parse the signature
        let signature_bytes = hex::decode(signature)
            .map_err(|_| anyhow!("Invalid signature format"))?;
        
        let signature = Signature::from_compact(&signature_bytes)
            .map_err(|_| anyhow!("Invalid signature"))?;

        // Create message hash from nonce
        let message_hash = Sha256::digest(nonce_response.nonce.as_bytes());
        let message = Message::from_slice(&message_hash)
            .map_err(|_| anyhow!("Failed to create message"))?;

        // Verify the signature
        match secp.verify_ecdsa(&message, &signature, &public_key) {
            Ok(_) => {
                // Mark this key as authenticated
                self.authenticated_keys.insert(pubkey.to_string(), current_time);
                // Remove the used nonce
                self.nonces.remove(pubkey);
                Ok(true)
            }
            Err(_) => Ok(false),
        }
    }

    /// Check if a public key is currently authenticated
    pub fn is_authenticated(&self, pubkey: &str) -> bool {
        if let Some(&timestamp) = self.authenticated_keys.get(pubkey) {
            let current_time = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs();
            
            // Authentication is valid for 1 hour
            current_time - timestamp < 3600
        } else {
            false
        }
    }

    /// Remove expired authentications
    pub fn cleanup_expired(&mut self) {
        let current_time = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        // Remove expired nonces (older than 5 minutes)
        self.nonces.retain(|_, nonce| current_time - nonce.timestamp <= 300);

        // Remove expired authentications (older than 1 hour)
        self.authenticated_keys.retain(|_, &mut timestamp| current_time - timestamp <= 3600);
    }
}