use std::sync::Arc;

use bitcoin::secp256k1::PublicKey;
use frost_secp256k1_tr::Identifier;
use serde::{Deserialize, Serialize};
use serde_with::{DisplayFromStr, serde_as};

use crate::{
    header_provider::{CombinedHeaderProvider, HeaderProvider},
    operator::rpc::{ConnectionManager, OperatorRpcError, SoAuthHeaderProvider, SparkRpcClient},
    session_store::SessionStore,
    signer::SparkSigner,
};

use super::OperatorError;

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct OperatorPoolConfig {
    coordinator_index: usize,
    operators: Vec<OperatorConfig>,
}

impl OperatorPoolConfig {
    pub fn new(
        coordinator_index: usize,
        operators: Vec<OperatorConfig>,
    ) -> Result<Self, OperatorError> {
        if coordinator_index >= operators.len() {
            return Err(OperatorError::InvalidCoordinatorIndex);
        }

        // Ensure the operator ids make sense and are in order.
        for (index, operator) in operators.iter().enumerate() {
            if operator.id != index {
                return Err(OperatorError::InvalidOperatorId);
            }
        }

        Ok(Self {
            coordinator_index,
            operators,
        })
    }

    /// Returns the coordinator operator.
    pub fn get_coordinator(&self) -> &OperatorConfig {
        self.operators.get(self.coordinator_index).unwrap()
    }

    /// Returns an iterator over all operators, including the coordinator.
    pub fn get_all_operators(&self) -> impl Iterator<Item = &OperatorConfig> {
        self.operators.iter()
    }

    /// Returns an iterator over all operators except the coordinator.
    pub fn get_non_coordinator_operators(&self) -> impl Iterator<Item = &OperatorConfig> {
        self.operators
            .iter()
            .filter(|op| op.id != self.coordinator_index)
    }

    /// Returns the operator at the given index.
    pub fn get_operator_by_id(&self, id: usize) -> Option<&OperatorConfig> {
        self.operators.get(id)
    }

    /// Sets the user agent for all operators in the pool.
    #[must_use]
    pub fn with_user_agent(mut self, user_agent: Option<String>) -> Self {
        for operator in &mut self.operators {
            operator.user_agent = user_agent.clone();
        }
        self
    }
}

#[serde_as]
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct OperatorConfig {
    pub id: usize,
    pub identifier: Identifier,
    pub address: String,
    pub ca_cert: Option<Vec<u8>>,
    #[serde_as(as = "DisplayFromStr")]
    pub identity_public_key: PublicKey,
    pub user_agent: Option<String>,
}

impl OperatorConfig {}

#[derive(Clone)]
pub struct Operator {
    pub client: SparkRpcClient,
    pub id: usize,
    pub identifier: Identifier,
    pub identity_public_key: PublicKey,
}

pub struct OperatorPool {
    coordinator_index: usize,
    operators: Vec<Operator>,
}

impl OperatorPool {
    pub async fn connect(
        config: &OperatorPoolConfig,
        connection_manager: Arc<dyn ConnectionManager>,
        session_store: Arc<dyn SessionStore>,
        spark_signer: Arc<dyn SparkSigner>,
        extra_header_provider: Option<Arc<dyn HeaderProvider>>,
    ) -> Result<Self, OperatorRpcError> {
        let mut operators = Vec::new();
        for operator in &config.operators {
            let transport = connection_manager.get_transport(operator).await?;
            let auth_provider = Arc::new(SoAuthHeaderProvider::new(
                transport.clone(),
                Arc::clone(&spark_signer),
                session_store.clone(),
                operator.identity_public_key,
            ));
            let header_provider: Arc<dyn HeaderProvider> = match &extra_header_provider {
                Some(extra) => Arc::new(CombinedHeaderProvider::new(vec![
                    auth_provider,
                    Arc::clone(extra),
                ])),
                None => auth_provider,
            };
            let client = SparkRpcClient::new(transport, header_provider, operator.id);
            operators.push(Operator {
                client,
                id: operator.id,
                identifier: operator.identifier,
                identity_public_key: operator.identity_public_key,
            });
        }

        Ok(Self {
            coordinator_index: config.coordinator_index,
            operators,
        })
    }
    /// Returns the coordinator operator.
    pub fn get_coordinator(&self) -> &Operator {
        self.operators.get(self.coordinator_index).unwrap()
    }

    /// Returns an iterator over all operators, including the coordinator.
    pub fn get_all_operators(&self) -> impl Iterator<Item = &Operator> {
        self.operators.iter()
    }

    /// Returns an iterator over all operators except the coordinator.
    pub fn get_non_coordinator_operators(&self) -> impl Iterator<Item = &Operator> {
        self.operators
            .iter()
            .filter(|op| op.id != self.coordinator_index)
    }

    /// Returns the operator at the given index.
    pub fn get_operator_by_id(&self, id: usize) -> Option<&Operator> {
        self.operators.get(id)
    }

    /// Returns the operator with the given identifier.
    pub fn get_operator_by_identifier(&self, identifier: &Identifier) -> Option<&Operator> {
        self.operators
            .iter()
            .find(|op| &op.identifier == identifier)
    }

    pub fn is_empty(&self) -> bool {
        self.operators.is_empty()
    }

    pub fn len(&self) -> usize {
        self.operators.len()
    }
}