use crate::config::AppConfig;
use crate::error::Result;
use crate::persistence::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,
}

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

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(RwLock::new(HashMap::new())),
        })
    }

    /// Finds a healthy service instance by name (e.g., "tushare").
    /// Returns the base_url.
    pub fn get_service_url(&self, service_name: &str) -> Option<String> {
        let registry = self.registry.read().unwrap();
        // TODO: Implement Round-Robin or check Last Heartbeat for health?
        // For now, return the first match.
        registry
            .values()
            .find(|entry| entry.registration.service_name == service_name)
            .map(|entry| entry.registration.base_url.clone())
    }

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

    pub fn get_provider_count(&self) -> usize {
        let registry = self.registry.read().unwrap();
        registry
            .values()
            .filter(|entry| {
                // Strict type checking using ServiceRole
                entry.registration.role == ServiceRole::DataProvider
            })
            .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())
}
