use crate::config::Config;
use crate::error::{HermesError, Result};
use crate::identity::resolve_profile;
use crate::network::message::HermesMessage;
use crate::network::node::{HermesNode, DATA_ALPN};
use crate::network::queue::MessageQueue;
use crate::store::consumer::push_file;
use crate::store::download::download_file_from_store;
use crate::store::provider::handle_data_push;
use crate::store::upload::upload_file_to_store;

use iroh::{EndpointAddr, PublicKey, SecretKey};
use std::collections::HashSet;
use std::sync::Arc;
use tokio::sync::Semaphore;
use tokio::time::{sleep, Duration};

/// Maximum concurrent incoming QUIC connections handled simultaneously.
const MAX_CONCURRENT_CONNECTIONS: usize = 128;

/// Maximum retries before a queued message is dropped permanently.
const MAX_RETRY_COUNT: u32 = 100;

#[derive(serde::Serialize)]
struct PayloadMeta {
    hash: String,
    encrypted_keys: Vec<u8>,
    ephemeral_pub: Vec<u8>,
    dh_nonce: Vec<u8>,
}

/// The central Hermes Client that orchestrates Iroh P2P networking and the HTTP Sync-Engine.
#[derive(Clone)]
pub struct Client {
    config: Arc<Config>,
    pub(crate) node: HermesNode,
    pub(crate) queue: MessageQueue,
    http: reqwest::Client,
}

impl Client {
    /// Initializes the Hermes Client, booting up the Iroh P2P Node and Sled Database.
    pub async fn new(config: Config) -> Result<Self> {
        let key_bytes = std::fs::read(&config.node_secret_key_path).map_err(|e| {
            HermesError::Config(format!(
                "Failed to read secret key from {:?}: {}",
                config.node_secret_key_path, e
            ))
        })?;
        let key_array: [u8; 32] = key_bytes.as_slice().try_into().map_err(|_| {
            HermesError::Config(format!(
                "Secret key at {:?} must be exactly 32 bytes, got {}",
                config.node_secret_key_path,
                key_bytes.len()
            ))
        })?;
        let secret_key = SecretKey::from_bytes(&key_array);
        let node = HermesNode::new(secret_key, &config.subnet_ids).await?;

        let queue = MessageQueue::new(&config.storage_directory)?;
        let http = reqwest::Client::new();

        tracing::info!(
            ss58 = %config.ss58_address,
            subnets = ?config.subnet_ids,
            "Hermes client initialized"
        );

        Ok(Self {
            config: Arc::new(config),
            node,
            queue,
            http,
        })
    }

    /// Resolves the destination's Iroh EndpointAddr from their AccountProfile.
    async fn resolve_dest_addr(
        &self,
        dest_ss58: &str,
    ) -> Result<(EndpointAddr, crate::identity::AccountProfile)> {
        let profile = resolve_profile(&self.config.rpc_url, dest_ss58).await?;

        let node_id_array: [u8; 32] =
            profile.node_id.as_slice().try_into().map_err(|_| {
                HermesError::Identity("NodeId maps to incorrect byte length".into())
            })?;
        let node_id = PublicKey::from_bytes(&node_id_array)
            .map_err(|e| HermesError::Identity(format!("Invalid NodeId bytes: {}", e)))?;

        Ok((EndpointAddr::from(node_id), profile))
    }

    /// Sends a HermesMessage, buffering to the offline queue on failure.
    async fn send_or_queue(
        &self,
        dest_ss58: &str,
        dest_addr: EndpointAddr,
        msg: HermesMessage,
        subnet_id: Option<u16>,
    ) {
        if let Err(e) = self
            .node
            .send_message(dest_addr.clone(), msg.clone(), subnet_id)
            .await
        {
            tracing::debug!(dest = %dest_ss58, error = %e, "Direct send failed, queuing for retry");
            if let Err(qe) = self.queue.push(dest_ss58, dest_addr, &msg, subnet_id) {
                tracing::warn!(dest = %dest_ss58, error = %qe, "Failed to enqueue message for offline retry");
            }
        }
    }

    /// Primary action: Securely encrypt and stream a file directly to the destination via QUIC.
    pub async fn send_file(&self, _dest_ss58: &str, _file_path: &str) -> Result<String> {
        Err(HermesError::Unimplemented(
            "AES-GCM encrypted direct P2P streaming is not yet implemented. \
             Use `send_file_unencrypted` for direct P2P or `send_file_unencrypted_to_store` for Sync-Engine.".into()
        ))
    }

    /// Pushes a file directly to the destination peer over QUIC (DATA_ALPN).
    /// No intermediary HTTP server — pure P2P streaming.
    pub async fn send_file_unencrypted(&self, dest_ss58: &str, file_path: &str) -> Result<String> {
        let (dest_addr, _profile) = self.resolve_dest_addr(dest_ss58).await?;

        let filename = push_file(
            &self.node.endpoint,
            dest_addr,
            &self.config.ss58_address,
            file_path,
        )
        .await?;

        Ok(filename)
    }

    /// Download and decrypt a payload using the DH-encrypted AES key material to a specific dir.
    pub async fn receive_file(
        &self,
        _hash: &str,
        _encrypted_keys: &[u8],
        _ephemeral_pub: &[u8],
        _dh_nonce: &[u8],
        _out_dir: &str,
    ) -> Result<String> {
        Err(HermesError::Unimplemented(
            "AES-GCM encrypted receiving is not yet implemented. \
             Use the data callback on `spawn_listener` for direct P2P, or `receive_file_unencrypted_from_store` for Sync-Engine.".into()
        ))
    }

    /// Pull-based unencrypted file download (future feature).
    pub async fn receive_file_unencrypted(&self, _hash: &str, _out_dir: &str) -> Result<String> {
        Err(HermesError::Unimplemented(
            "Pull-based direct P2P receiving is not yet implemented. \
             Use the data callback on `spawn_listener` for push-based P2P receiving, \
             or `receive_file_unencrypted_from_store` for Sync-Engine."
                .into(),
        ))
    }

    /// Uploads a file to the Hippius Sync-Engine, sending the keys securely directly to the destination via Iroh.
    pub async fn send_file_to_store(&self, _dest_ss58: &str, _file_path: &str) -> Result<String> {
        Err(HermesError::Unimplemented(
            "Streaming AES-GCM encryption for Sync-Engine workloads is not yet fully implemented. \
             Use `send_file_unencrypted_to_store` for now."
                .into(),
        ))
    }

    /// Uploads an unencrypted file to the Hippius Sync-Engine.
    pub async fn send_file_unencrypted_to_store(
        &self,
        dest_ss58: &str,
        file_path: &str,
    ) -> Result<String> {
        let (dest_addr, _) = self.resolve_dest_addr(dest_ss58).await?;

        let hash = upload_file_to_store(
            &self.http,
            &self.config.sync_engine_url,
            &self.config.api_token,
            file_path,
        )
        .await?;

        let meta = PayloadMeta {
            hash: hash.clone(),
            encrypted_keys: vec![],
            ephemeral_pub: vec![],
            dh_nonce: vec![],
        };

        let msg = HermesMessage {
            action: "process_data_unencrypted_store".into(),
            sender_ss58: self.config.ss58_address.clone(),
            payload: serde_json::to_vec(&meta)?,
        };

        self.send_or_queue(dest_ss58, dest_addr, msg, None).await;
        Ok(hash)
    }

    /// Downloads and decrypts a file from the Hippius Sync-Engine to a specific dir.
    pub async fn receive_file_from_store(
        &self,
        _hash: &str,
        _encrypted_keys: &[u8],
        _ephemeral_pub: &[u8],
        _dh_nonce: &[u8],
        _out_dir: &str,
    ) -> Result<String> {
        Err(HermesError::Unimplemented(
            "Streaming AES-GCM decryption for Sync-Engine workloads is not yet fully implemented. \
             Use `receive_file_unencrypted_from_store` for now."
                .into(),
        ))
    }

    /// Downloads an unencrypted file directly from the Hippius Sync-Engine.
    pub async fn receive_file_unencrypted_from_store(
        &self,
        hash: &str,
        out_dir: &str,
    ) -> Result<String> {
        let path =
            download_file_from_store(&self.http, &self.config.sync_engine_url, hash, out_dir)
                .await?;
        Ok(path)
    }

    /// Spawns an async loop to drain the Sled queue and retry sending offline messages.
    /// Groups by destination to avoid hammering offline peers. Messages exceeding
    /// MAX_RETRY_COUNT are dropped. Failed destinations are skipped for the remainder
    /// of each sweep to prevent head-of-line blocking.
    pub fn spawn_retry_worker(&self) {
        let client = self.clone();
        tokio::spawn(async move {
            let retry_semaphore = Arc::new(Semaphore::new(50));

            loop {
                // Drain everything currently in the queue
                let mut items = Vec::new();
                while let Ok(Some(item)) = client.queue.pop_next() {
                    items.push(item);
                }

                if !items.is_empty() {
                    let failed_dests = Arc::new(tokio::sync::Mutex::new(HashSet::<String>::new()));
                    let mut handles = Vec::new();

                    for (dest_ss58, dest_addr, msg, retry_count, subnet_id) in items {
                        // Drop messages that have exceeded the retry limit
                        if retry_count >= MAX_RETRY_COUNT {
                            tracing::warn!(dest = %dest_ss58, retries = retry_count, "Dropping message after max retries");
                            continue;
                        }

                        let c = client.clone();
                        let fd = failed_dests.clone();
                        let permit = retry_semaphore.clone().acquire_owned().await.unwrap();

                        handles.push(tokio::spawn(async move {
                            let _permit = permit;

                            // Skip destinations already known to be offline this sweep
                            {
                                let dests = fd.lock().await;
                                if dests.contains(&dest_ss58) {
                                    if let Err(e) = c.queue.push_retry(&dest_ss58, dest_addr, &msg, retry_count, subnet_id) {
                                        tracing::warn!(dest = %dest_ss58, error = %e, "Failed to re-queue skipped message");
                                    }
                                    return;
                                }
                            }

                            if let Err(e) = c.node.send_message(dest_addr.clone(), msg.clone(), subnet_id).await {
                                tracing::debug!(dest = %dest_ss58, error = %e, retry = retry_count + 1, "Retry send failed");
                                {
                                    let mut dests = fd.lock().await;
                                    dests.insert(dest_ss58.clone());
                                }
                                if let Err(qe) = c.queue.push_retry(&dest_ss58, dest_addr, &msg, retry_count + 1, subnet_id) {
                                    tracing::warn!(dest = %dest_ss58, error = %qe, "Failed to re-queue message after retry failure");
                                }
                            }
                        }));
                    }

                    // Wait for all retry tasks before next sweep
                    for handle in handles {
                        if let Err(e) = handle.await {
                            tracing::warn!(error = %e, "Retry task panicked");
                        }
                    }
                }

                // Sleep before initiating the next full sweep
                sleep(Duration::from_secs(60)).await;
            }
        });
    }

    /// Spawns an async loop to accept incoming Iroh QUIC streams and route them.
    ///
    /// - **Control ALPN connections** are routed to `control_callback` after identity verification.
    /// - **Data ALPN connections** are routed to `data_callback` (if provided) for direct P2P file reception.
    ///
    /// If `data_callback` is `None`, incoming data connections are silently dropped.
    pub fn spawn_listener<F, D>(&self, control_callback: F, data_callback: Option<D>)
    where
        F: Fn(HermesMessage) + Send + Sync + 'static,
        D: Fn(String, String, String, u64) + Send + Sync + 'static,
    {
        let endpoint = self.node.endpoint.clone();
        let registered_alpns = self.node.registered_alpns.clone();
        let rpc_url = self.config.rpc_url.clone();
        let storage_dir = self.config.storage_directory.clone();
        let control_callback = Arc::new(control_callback);
        let data_callback = data_callback.map(Arc::new);
        let semaphore = Arc::new(Semaphore::new(MAX_CONCURRENT_CONNECTIONS));

        tracing::info!("Spawning QUIC listener");

        tokio::spawn(async move {
            while let Some(incoming) = endpoint.accept().await {
                let cb = control_callback.clone();
                let dcb = data_callback.clone();
                let rpc = rpc_url.clone();
                let alpns = registered_alpns.clone();
                let storage = storage_dir.clone();
                let permit = match semaphore.clone().try_acquire_owned() {
                    Ok(p) => p,
                    Err(_) => continue, // At connection limit, drop this one
                };

                tokio::spawn(async move {
                    let _permit = permit; // Held until this handler completes

                    let connection = match incoming.await {
                        Ok(c) => c,
                        Err(_) => return,
                    };

                    // Defense-in-depth: reject connections whose ALPN doesn't match
                    let alpn = connection.alpn();
                    if !alpns.iter().any(|a| a == alpn) {
                        return;
                    }

                    // Route based on ALPN
                    if alpn == DATA_ALPN {
                        // Direct P2P data transfer
                        let dcb = match dcb {
                            Some(ref cb) => cb.clone(),
                            None => return, // No data callback, silently drop
                        };

                        // Verify sender identity from the TLS-authenticated remote NodeId
                        let remote_pub = connection.remote_id();

                        if let Ok((sender_ss58, filename, local_path, file_size)) =
                            handle_data_push(connection, &storage).await
                        {
                            // Verify: the sender's claimed SS58 must map to the TLS-authenticated NodeId
                            let verified = match resolve_profile(&rpc, &sender_ss58).await {
                                Ok(profile) => {
                                    let expected: [u8; 32] =
                                        match profile.node_id.as_slice().try_into() {
                                            Ok(arr) => arr,
                                            Err(_) => return,
                                        };
                                    match PublicKey::from_bytes(&expected) {
                                        Ok(expected_pub) => expected_pub == remote_pub,
                                        Err(_) => false,
                                    }
                                }
                                Err(_) => false,
                            };

                            if verified {
                                dcb(sender_ss58, filename, local_path, file_size);
                            }
                        }
                    } else {
                        // Control plane message
                        let remote_pub = connection.remote_id();

                        let (_send, mut recv) = match connection.accept_bi().await {
                            Ok(bi) => bi,
                            Err(_) => return,
                        };

                        if let Ok(buf) = recv.read_to_end(10_000_000).await {
                            if let Ok(msg) = serde_json::from_slice::<HermesMessage>(&buf) {
                                let verified = match resolve_profile(&rpc, &msg.sender_ss58).await {
                                    Ok(profile) => {
                                        let expected: [u8; 32] =
                                            match profile.node_id.as_slice().try_into() {
                                                Ok(arr) => arr,
                                                Err(_) => return,
                                            };
                                        match PublicKey::from_bytes(&expected) {
                                            Ok(expected_pub) => expected_pub == remote_pub,
                                            Err(_) => false,
                                        }
                                    }
                                    Err(_) => false,
                                };

                                if verified {
                                    cb(msg);
                                }
                            }
                        }
                    }
                });
            }
        });
    }
}
