use crate::config::AppConfig;
use crate::error::Result;
use common_contracts::persistence_client::PersistenceClient;
use async_nats::Client as NatsClient;
use common_contracts::registry::{ServiceRegistration, ServiceRole};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Instant;
use tokio::time::{Duration, sleep};
use tracing::{info, warn};

#[derive(Clone, Debug)]
pub struct RegistryEntry {
    pub registration: ServiceRegistration,
    pub last_heartbeat: Instant,
}

/// Encapsulates the service registry logic to make it testable without NATS.
#[derive(Debug, Default)]
pub struct ServiceRegistry {
    services: RwLock<HashMap<String, RegistryEntry>>,
}

impl ServiceRegistry {
    pub fn new() -> Self {
        Self {
            services: RwLock::new(HashMap::new()),
        }
    }

    pub fn register(&self, registration: ServiceRegistration) {
        let mut services = self.services.write().unwrap();
        services.insert(
            registration.service_id.clone(),
            RegistryEntry {
                registration,
                last_heartbeat: Instant::now(),
            },
        );
    }

    pub fn deregister(&self, service_id: &str) -> bool {
        let mut services = self.services.write().unwrap();
        services.remove(service_id).is_some()
    }

    pub fn update_heartbeat(&self, service_id: &str) -> bool {
        let mut services = self.services.write().unwrap();
        if let Some(entry) = services.get_mut(service_id) {
            entry.last_heartbeat = Instant::now();
            true
        } else {
            false
        }
    }

    pub fn get_service_url(&self, service_name: &str) -> Option<String> {
        let services = self.services.read().unwrap();
        services
            .values()
            .find(|entry| entry.registration.service_name == service_name)
            .map(|entry| entry.registration.base_url.clone())
    }

    pub fn get_all_services(&self) -> Vec<(String, String)> {
        let services = self.services.read().unwrap();
        services
            .values()
            .map(|entry| {
                (
                    entry.registration.service_id.clone(),
                    entry.registration.base_url.clone(),
                )
            })
            .collect()
    }

    #[allow(dead_code)]
    pub fn get_provider_count(&self) -> usize {
        let services = self.services.read().unwrap();
        services
            .values()
            .filter(|entry| entry.registration.role == ServiceRole::DataProvider)
            .count()
    }
    
    pub fn get_entries(&self) -> Vec<RegistryEntry> {
        let services = self.services.read().unwrap();
        services.values().cloned().collect()
    }
}

#[derive(Clone)]
pub struct AppState {
    pub config: Arc<AppConfig>,
    pub nats_client: NatsClient,
    pub persistence_client: PersistenceClient,
    // Replaced Arc<RwLock<HashMap>> with Arc<ServiceRegistry>
    pub registry: Arc<ServiceRegistry>,
}

impl AppState {
    pub async fn new(config: AppConfig) -> Result<Self> {
        let nats_client = connect_nats_with_retry(&config.nats_addr).await?;

        let persistence_client =
            PersistenceClient::new(config.data_persistence_service_url.clone());

        Ok(Self {
            config: Arc::new(config),
            nats_client,
            persistence_client,
            registry: Arc::new(ServiceRegistry::new()),
        })
    }

    // Delegate methods to registry
    pub fn get_service_url(&self, service_name: &str) -> Option<String> {
        self.registry.get_service_url(service_name)
    }

    pub fn get_all_services(&self) -> Vec<(String, String)> {
        self.registry.get_all_services()
    }

    #[allow(dead_code)]
    pub fn get_provider_count(&self) -> usize {
        self.registry.get_provider_count()
    }
}

const MAX_NATS_CONNECT_ATTEMPTS: usize = 30;
const INITIAL_BACKOFF_MS: u64 = 200;
const MAX_BACKOFF_MS: u64 = 2_000;

async fn connect_nats_with_retry(nats_addr: &str) -> Result<NatsClient> {
    let mut backoff_ms = INITIAL_BACKOFF_MS;
    let mut last_err: Option<async_nats::error::Error<async_nats::ConnectErrorKind>> = None;

    for attempt in 1..=MAX_NATS_CONNECT_ATTEMPTS {
        match async_nats::connect(nats_addr).await {
            Ok(client) => {
                info!("Connected to NATS at {} (attempt {})", nats_addr, attempt);
                return Ok(client);
            }
            Err(err) => {
                warn!(
                    "Failed to connect to NATS at {} (attempt {}): {}",
                    nats_addr, attempt, err
                );
                last_err = Some(err);
                let delay = backoff_ms.min(MAX_BACKOFF_MS);
                sleep(Duration::from_millis(delay)).await;
                backoff_ms = backoff_ms.saturating_mul(2);
            }
        }
    }

    // Exhausted all attempts; return the last error
    Err(last_err
        .expect("last_err must be set after failed attempts")
        .into())
}
