use async_std::task;
use futures::prelude::*;
use libp2p::{
    development_transport,
    identity,
    kad::{
        record::{Key as KadKey, Record},
        Kademlia, KademliaEvent, PutRecordOk, GetRecordOk, Quorum, 
        record::store::MemoryStore,
    },
    swarm::{NetworkBehaviour, Swarm, SwarmEvent},
    Multiaddr, PeerId,
};
use serde::{Deserialize, Serialize};
use serde_json;
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::error::Error;

// ---------------------------
// Data Structures
// ---------------------------

/// A stored immutable document on the network.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Document {
    /// For agents, this is "agent_id:version"; for documents it can be "doc_id:version"
    pub agent_id: String,
    /// Document id (e.g. "public_key" or "recovery_code")
    pub document_id: String,
    /// The actual value (e.g. the public key as hex or the recovery code)
    pub value: String,
    /// The signature (as hex) over the value (or over a canonical record)
    pub signature: String,
}

/// An Agent holds the identity information in our system.
pub struct Agent {
    /// The agent id is in the form "agent_id:version"
    pub id: String,
    /// The current keypair
    pub keypair: ed25519_dalek::Keypair,
    /// A large recovery code (a long random string)
    pub recovery_code: String,
    /// The recovery code’s signature (signed by the agent’s private key)
    pub recovery_signature: ed25519_dalek::Signature,
}

// ---------------------------
// libp2p Behaviour
// ---------------------------

/// Our network behaviour currently consists solely of a Kademlia DHT.
#[derive(NetworkBehaviour)]
#[behaviour(event_process = true)]
struct MyBehaviour {
    kad: Kademlia<MemoryStore>,
}

impl MyBehaviour {
    pub fn new(peer_id: PeerId) -> Self {
        let store = MemoryStore::new(peer_id);
        let kad = Kademlia::new(peer_id, store);
        MyBehaviour { kad }
    }
}

// For demonstration, we simply print Kademlia events.
impl libp2p::swarm::NetworkBehaviourEventProcess<KademliaEvent> for MyBehaviour {
    fn inject_event(&mut self, event: KademliaEvent) {
        println!("Kademlia event: {:?}", event);
    }
}

/// A wrapper that uses libp2p’s Kademlia DHT to store/retrieve documents.
pub struct NetworkKVStore {
    pub swarm: Swarm<MyBehaviour>,
}

impl NetworkKVStore {
    /// Create a new NetworkKVStore with a libp2p swarm.
    pub async fn new() -> Self {
        // Generate a libp2p identity for the node.
        let local_key = identity::Keypair::generate_ed25519();
        let local_peer_id = PeerId::from(local_key.public());
        println!("Local peer id: {:?}", local_peer_id);

        // Build a transport (TCP + noise, multiplexed).
        let transport = development_transport(local_key.clone()).await.unwrap();

        // Create our Kademlia behaviour.
        let behaviour = MyBehaviour::new(local_peer_id);
        let mut swarm = Swarm::new(transport, behaviour, local_peer_id);

        // Start listening on an ephemeral TCP port.
        let addr: Multiaddr = "/ip4/0.0.0.0/tcp/0".parse().unwrap();
        Swarm::listen_on(&mut swarm, addr).unwrap();

        NetworkKVStore { swarm }
    }

    /// Writes a document to the DHT. We serialize the Document into JSON bytes.
    pub async fn write_document(&mut self, doc: Document) -> Result<(), Box<dyn Error>> {
        // Use a key prefix "doc:" so that document ids don’t conflict.
        let key = format!("doc:{}", doc.document_id);
        let kad_key = KadKey::new(&key);
        let value = serde_json::to_vec(&doc)?;
        let record = Record {
            key: kad_key,
            value,
            publisher: None,
            expires: None,
        };

        let query_id = self.swarm.behaviour_mut().kad.put_record(record, Quorum::One)?;
        loop {
            match self.swarm.next().await {
                Some(SwarmEvent::Behaviour(KademliaEvent::OutboundQueryCompleted { id, result, .. })) if id == query_id => {
                    match result {
                        libp2p::kad::QueryResult::PutRecord(Ok(PutRecordOk { key, .. })) => {
                            println!("Successfully put record with key: {:?}", key);
                            return Ok(());
                        }
                        libp2p::kad::QueryResult::PutRecord(Err(err)) => {
                            println!("Failed to put record: {:?}", err);
                            return Err("Failed to put record".into());
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        }
    }

    /// Retrieves a document from the DHT given its document_id.
    pub async fn get_document(&mut self, document_id: &str) -> Result<Document, Box<dyn Error>> {
        let key = format!("doc:{}", document_id);
        let kad_key = KadKey::new(&key);
        let query_id = self.swarm.behaviour_mut().kad.get_record(kad_key, Quorum::One);
        loop {
            match self.swarm.next().await {
                Some(SwarmEvent::Behaviour(KademliaEvent::OutboundQueryCompleted { id, result, .. })) if id == query_id => {
                    match result {
                        libp2p::kad::QueryResult::GetRecord(Ok(GetRecordOk { records, .. })) => {
                            if let Some(record_entry) = records.first() {
                                let doc: Document = serde_json::from_slice(&record_entry.record.value)?;
                                return Ok(doc);
                            }
                        }
                        libp2p::kad::QueryResult::GetRecord(Err(err)) => {
                            println!("Error in get_record: {:?}", err);
                            return Err("Failed to get record".into());
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        }
    }
}

// ---------------------------
// Helper Functions for Agent & Signing
// ---------------------------
use ed25519_dalek::{Keypair, PublicKey, Signature, Signer};
use rand::rngs::OsRng;
use rand::Rng;

/// Generates a large recovery code as a random 64-character hexadecimal string.
fn generate_large_recovery_code() -> String {
    let mut rng = rand::thread_rng();
    let code: [u8; 32] = rng.gen();
    hex::encode(code)
}

/// Produces a placeholder signature for a new public key using the recovery code.
fn new_public_key_signature_placeholder(new_public_key: &PublicKey, recovery_code: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(new_public_key.to_bytes());
    hasher.update(recovery_code.as_bytes());
    let result = hasher.finalize();
    hex::encode(result)
}

/// Creates a new agent:
/// 1. Generates a keypair and a recovery code.
/// 2. Signs the recovery code with the private key.
/// 3. Writes two documents: one for the public key and one for the recovery code.
pub async fn create_agent(nkv: &mut NetworkKVStore) -> Agent {
    let mut csprng = OsRng {};
    let keypair: Keypair = Keypair::generate(&mut csprng);
    let pub_key_bytes = keypair.public.to_bytes();
    let agent_id = format!("{}:1", hex::encode(&pub_key_bytes));

    let recovery_code = generate_large_recovery_code();
    let recovery_signature = keypair.sign(recovery_code.as_bytes());

    // Create and write the public key document.
    let pub_doc = Document {
        agent_id: agent_id.clone(),
        document_id: format!("{}:public_key", agent_id),
        value: hex::encode(&pub_key_bytes),
        signature: hex::encode(keypair.sign(&pub_key_bytes).to_bytes()),
    };
    nkv.write_document(pub_doc).await.unwrap();

    // Create and write the recovery code document.
    let rec_doc = Document {
        agent_id: agent_id.clone(),
        document_id: format!("{}:recovery_code", agent_id),
        value: recovery_code.clone(),
        signature: hex::encode(recovery_signature.to_bytes()),
    };
    nkv.write_document(rec_doc).await.unwrap();

    Agent { id: agent_id, keypair, recovery_code, recovery_signature }
}

/// Updates the agent’s public key using the recovery code.
/// It verifies the stored recovery code, increments the version,
/// and writes a new public key document.
pub async fn update_agent(
    agent: &Agent,
    new_public_key: PublicKey,
    nkv: &mut NetworkKVStore,
    provided_recovery_code: &str,
) -> bool {
    let recovery_doc_id = format!("{}:recovery_code", agent.id);
    let stored_recovery = nkv.get_document(&recovery_doc_id).await;
    if stored_recovery.is_err() {
        println!("No recovery document found.");
        return false;
    }
    let stored_doc = stored_recovery.unwrap();
    if stored_doc.value != provided_recovery_code {
        println!("Recovery code mismatch.");
        return false;
    }

    // Increment the version.
    let parts: Vec<&str> = agent.id.split(':').collect();
    if parts.len() != 2 {
        println!("Invalid agent id format.");
        return false;
    }
    let version: u64 = parts[1].parse().unwrap_or(1);
    let new_version = version + 1;
    let new_pub_key_bytes = new_public_key.to_bytes();
    let new_agent_id = format!("{}:{}", hex::encode(&new_pub_key_bytes), new_version);
    let new_pk_sig = new_public_key_signature_placeholder(&new_public_key, provided_recovery_code);

    let pub_doc = Document {
        agent_id: new_agent_id.clone(),
        document_id: format!("{}:public_key", new_agent_id),
        value: hex::encode(&new_pub_key_bytes),
        signature: new_pk_sig,
    };

    nkv.write_document(pub_doc).await.is_ok()
}

/// Retrieves the signature for a document.
pub async fn get_signature(nkv: &mut NetworkKVStore, document_id: &str) -> Option<String> {
    nkv.get_document(document_id).await.ok().map(|doc| doc.signature)
}

/// Retrieves the public key for an agent.
pub async fn get_public_key(nkv: &mut NetworkKVStore, agent_doc_id: &str) -> Option<String> {
    nkv.get_document(agent_doc_id).await.ok().map(|doc| doc.value)
}

// ---------------------------
// Main function to test the network integration
// ---------------------------
#[async_std::main]
async fn main() {
    // Create the network KV store (this sets up our libp2p swarm with Kademlia).
    let mut nkv = NetworkKVStore::new().await;

    // Run the swarm in a background task so it can process events.
    task::spawn(async move {
        loop {
            if let Some(event) = nkv.swarm.next().await {
                println!("Swarm event: {:?}", event);
   
