use std::collections::HashMap;
use std::fs;
use std::path::Path;

use ed25519_dalek::{SigningKey, VerifyingKey};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use x25519_dalek::EphemeralSecret;

use backchannel_common::crypto::identity;

use crate::error::{ClientError, Result};

/// On-disk representation of the keystore.
#[derive(Serialize, Deserialize)]
struct KeystoreFile {
    /// Base64-encoded Ed25519 signing key (32 bytes).
    signing_key_b64: String,
}

/// Holds cryptographic material for the local client.
pub struct Keystore {
    pub signing_key:   SigningKey,
    pub verifying_key: VerifyingKey,

    /// Derived DM session keys cached in memory: peer_user_id → 32-byte key.
    /// Populated after completing an X25519 key exchange.
    pub dm_keys: HashMap<Uuid, [u8; 32]>,

    /// Ephemeral secrets awaiting the remote party's public key.
    /// Keyed by the peer's user UUID. Consumed on exchange completion.
    pub pending_ecdh: HashMap<Uuid, EphemeralSecret>,
}

impl Keystore {
    /// Load an existing keystore from `path`, or generate a fresh one and save it.
    pub fn load_or_create(path: &str) -> Result<Self> {
        if Path::new(path).exists() {
            let contents = fs::read_to_string(path)?;
            let file: KeystoreFile =
                serde_json::from_str(&contents).map_err(ClientError::Json)?;
            let signing_key   = identity::b64_to_signing_key(&file.signing_key_b64)?;
            let verifying_key = signing_key.verifying_key();
            tracing::debug!("Loaded keystore from {}", path);
            Ok(Self::from_signing_key(signing_key, verifying_key))
        } else {
            let (signing_key, verifying_key) = identity::generate_keypair();
            let file = KeystoreFile {
                signing_key_b64: identity::signing_key_to_b64(&signing_key),
            };
            if let Some(parent) = Path::new(path).parent() {
                fs::create_dir_all(parent)?;
            }
            fs::write(path, serde_json::to_string_pretty(&file)?)?;
            tracing::info!("Generated new identity keypair, saved to {}", path);
            Ok(Self::from_signing_key(signing_key, verifying_key))
        }
    }

    fn from_signing_key(signing_key: SigningKey, verifying_key: VerifyingKey) -> Self {
        Self {
            signing_key,
            verifying_key,
            dm_keys:      HashMap::new(),
            pending_ecdh: HashMap::new(),
        }
    }

    /// Base64-encoded Ed25519 public key, suitable for sending to the server.
    pub fn identity_pubkey_b64(&self) -> String {
        identity::pubkey_to_b64(&self.verifying_key)
    }

    /// Store an `EphemeralSecret` while waiting for the remote peer's response.
    pub fn store_pending_ecdh(&mut self, peer_id: Uuid, secret: EphemeralSecret) {
        self.pending_ecdh.insert(peer_id, secret);
    }

    /// Retrieve and remove a pending ephemeral secret (consumed by ECDH).
    pub fn take_pending_ecdh(&mut self, peer_id: Uuid) -> Option<EphemeralSecret> {
        self.pending_ecdh.remove(&peer_id)
    }

    pub fn cache_dm_key(&mut self, peer_id: Uuid, key: [u8; 32]) {
        self.dm_keys.insert(peer_id, key);
    }

    pub fn get_dm_key(&self, peer_id: &Uuid) -> Option<&[u8; 32]> {
        self.dm_keys.get(peer_id)
    }
}