//! Shared IAM auth helpers for service-side gRPC/HTTP handlers.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};

use http::header::AUTHORIZATION;
use http::HeaderMap;
use iam_client::client::IamClientConfig;
use iam_client::IamClient;
use iam_types::{Principal, PrincipalKind, Resource, Scope};
use tonic::metadata::MetadataMap;
use tonic::{Request, Status};
use tracing::{debug, warn};

const PHOTON_AUTH_TOKEN_HEADER: &str = "x-photon-auth-token";
const DEFAULT_TOKEN_CACHE_TTL_MS: u64 = 5_000;
const DEFAULT_AUTHZ_CACHE_TTL_MS: u64 = 3_000;
const AUTH_CONNECT_RETRY_ATTEMPTS: usize = 6;
const AUTH_CONNECT_INITIAL_BACKOFF: Duration = Duration::from_millis(500);
const AUTH_CONNECT_MAX_BACKOFF: Duration = Duration::from_secs(5);

#[derive(Debug, Clone)]
struct CacheEntry<T> {
    value: T,
    expires_at: Instant,
}

/// Tenant context extracted from authenticated token.
#[derive(Debug, Clone)]
pub struct TenantContext {
    pub org_id: String,
    pub project_id: String,
    pub principal_id: String,
    pub principal_name: String,
    pub principal_kind: PrincipalKind,
    pub node_id: Option<String>,
}

/// Authentication service that validates tokens and extracts tenant context.
#[derive(Clone)]
pub struct AuthService {
    iam_client: Arc<IamClient>,
    token_cache: Arc<RwLock<HashMap<String, CacheEntry<TenantContext>>>>,
    authz_cache: Arc<RwLock<HashMap<String, CacheEntry<bool>>>>,
    token_cache_ttl: Duration,
    authz_cache_ttl: Duration,
}

impl AuthService {
    /// Create a new authentication service.
    pub async fn new(iam_endpoint: &str) -> Result<Self, String> {
        let mut config = IamClientConfig::new(iam_endpoint).with_timeout(5000);
        let tls_override = resolve_tls_override(iam_endpoint);
        if let Some(enable_tls) = tls_override {
            if !enable_tls {
                config = config.without_tls();
            }
        } else if iam_endpoint.starts_with("http://") {
            config = config.without_tls();
        } else if !iam_endpoint.starts_with("https://") {
            // Backward-compatible default: plaintext for scheme-less endpoints.
            // Use https:// or PHOTON_IAM_TLS=true to enable TLS.
            warn!(
                "IAM endpoint has no scheme; defaulting to plaintext. Use https:// or PHOTON_IAM_TLS=true"
            );
            config = config.without_tls();
        }

        let iam_client = connect_iam_with_retry(config).await?;

        Ok(Self {
            iam_client: Arc::new(iam_client),
            token_cache: Arc::new(RwLock::new(HashMap::new())),
            authz_cache: Arc::new(RwLock::new(HashMap::new())),
            token_cache_ttl: cache_ttl_from_env(
                "PHOTON_AUTH_TOKEN_CACHE_TTL_MS",
                DEFAULT_TOKEN_CACHE_TTL_MS,
            ),
            authz_cache_ttl: cache_ttl_from_env(
                "PHOTON_AUTH_AUTHZ_CACHE_TTL_MS",
                DEFAULT_AUTHZ_CACHE_TTL_MS,
            ),
        })
    }

    /// Authenticate a gRPC request and return tenant context.
    pub async fn authenticate_request<T>(
        &self,
        request: &Request<T>,
    ) -> Result<TenantContext, Status> {
        let token = extract_token_from_metadata(request.metadata())?;
        self.authenticate_token(&token).await
    }

    /// Authenticate an HTTP request using headers.
    pub async fn authenticate_headers(&self, headers: &HeaderMap) -> Result<TenantContext, Status> {
        let token = extract_token_from_headers(headers)?;
        self.authenticate_token(&token).await
    }

    /// Authorize an action against a resource using IAM.
    pub async fn authorize(
        &self,
        tenant: &TenantContext,
        action: &str,
        resource: &Resource,
    ) -> Result<(), Status> {
        let mut principal = match tenant.principal_kind {
            PrincipalKind::User => {
                Principal::new_user(&tenant.principal_id, &tenant.principal_name)
            }
            PrincipalKind::ServiceAccount => Principal::new_service_account(
                &tenant.principal_id,
                &tenant.principal_name,
                &tenant.org_id,
                &tenant.project_id,
            ),
            PrincipalKind::Group => {
                Principal::new_group(&tenant.principal_id, &tenant.principal_name)
            }
        };

        principal.org_id = Some(tenant.org_id.clone());
        principal.project_id = Some(tenant.project_id.clone());
        principal.node_id = tenant.node_id.clone();

        let allowed = self.authorize_cached(&principal, action, resource).await?;

        if allowed {
            Ok(())
        } else {
            Err(Status::permission_denied("IAM policy denied the request"))
        }
    }

    async fn authenticate_token(&self, token: &str) -> Result<TenantContext, Status> {
        if let Some(cached) = self.cached_token(token).await {
            return Ok(cached);
        }

        let claims = self.iam_client.validate_token(token).await.map_err(|e| {
            warn!("Token validation failed: {}", e);
            Status::unauthenticated(format!("Invalid token: {}", e))
        })?;

        let org_id = claims
            .org_id
            .clone()
            .or_else(|| match &claims.scope {
                Scope::Org { id } => Some(id.clone()),
                Scope::Project { org_id, .. } => Some(org_id.clone()),
                Scope::Resource { org_id, .. } => Some(org_id.clone()),
                Scope::System => None,
            })
            .ok_or_else(|| {
                warn!("Token missing org_id");
                Status::unauthenticated("Token missing org_id")
            })?;

        let project_id = claims
            .project_id
            .clone()
            .or_else(|| match &claims.scope {
                Scope::Project { id, .. } => Some(id.clone()),
                Scope::Resource { project_id, .. } => Some(project_id.clone()),
                _ => None,
            })
            .ok_or_else(|| {
                warn!("Token missing project_id");
                Status::unauthenticated("Token missing project_id")
            })?;

        debug!(
            "Authenticated request: org_id={}, project_id={}, principal={}",
            org_id, project_id, claims.principal_id
        );

        let tenant = TenantContext {
            org_id,
            project_id,
            principal_id: claims.principal_id,
            principal_name: claims.principal_name,
            principal_kind: claims.principal_kind,
            node_id: claims.node_id,
        };
        self.store_token_cache(token, &tenant).await;
        Ok(tenant)
    }

    async fn authorize_cached(
        &self,
        principal: &Principal,
        action: &str,
        resource: &Resource,
    ) -> Result<bool, Status> {
        let cache_key = authz_cache_key(principal, action, resource).ok();
        if let Some(cache_key) = cache_key.as_deref() {
            if let Some(cached) = self.cached_authz(cache_key).await {
                return Ok(cached);
            }
        }

        let allowed = self
            .iam_client
            .authorize(principal, action, resource)
            .await
            .map_err(|e| Status::unavailable(format!("IAM authorization failed: {}", e)))?;

        if let Some(cache_key) = cache_key.as_deref() {
            self.store_authz_cache(cache_key, allowed).await;
        }
        Ok(allowed)
    }

    async fn cached_token(&self, token: &str) -> Option<TenantContext> {
        let cache = self.token_cache.read().ok()?;
        let entry = cache.get(token)?;
        if entry.expires_at > Instant::now() {
            return Some(entry.value.clone());
        }
        drop(cache);
        self.token_cache.write().ok()?.remove(token);
        None
    }

    async fn store_token_cache(&self, token: &str, tenant: &TenantContext) {
        if self.token_cache_ttl.is_zero() {
            return;
        }

        let mut cache = match self.token_cache.write() {
            Ok(cache) => cache,
            Err(_) => return,
        };
        prune_expired(&mut cache);
        cache.insert(
            token.to_string(),
            CacheEntry {
                value: tenant.clone(),
                expires_at: Instant::now() + self.token_cache_ttl,
            },
        );
    }

    async fn cached_authz(&self, cache_key: &str) -> Option<bool> {
        let cache = self.authz_cache.read().ok()?;
        let entry = cache.get(cache_key)?;
        if entry.expires_at > Instant::now() {
            return Some(entry.value);
        }
        drop(cache);
        self.authz_cache.write().ok()?.remove(cache_key);
        None
    }

    async fn store_authz_cache(&self, cache_key: &str, allowed: bool) {
        if self.authz_cache_ttl.is_zero() {
            return;
        }

        let mut cache = match self.authz_cache.write() {
            Ok(cache) => cache,
            Err(_) => return,
        };
        prune_expired(&mut cache);
        cache.insert(
            cache_key.to_string(),
            CacheEntry {
                value: allowed,
                expires_at: Instant::now() + self.authz_cache_ttl,
            },
        );
    }
}

async fn connect_iam_with_retry(config: IamClientConfig) -> Result<IamClient, String> {
    let mut last_error = None;
    for attempt in 0..AUTH_CONNECT_RETRY_ATTEMPTS {
        match IamClient::connect(config.clone()).await {
            Ok(client) => return Ok(client),
            Err(err)
                if attempt + 1 < AUTH_CONNECT_RETRY_ATTEMPTS
                    && retryable_connect_error(&err.to_string()) =>
            {
                let delay = auth_connect_retry_delay(attempt);
                warn!(
                    attempt = attempt + 1,
                    retry_after_ms = delay.as_millis() as u64,
                    error = %err,
                    "retrying IAM auth service bootstrap connection"
                );
                last_error = Some(err.to_string());
                tokio::time::sleep(delay).await;
            }
            Err(err) => {
                return Err(format!("Failed to connect to IAM server: {}", err));
            }
        }
    }

    Err(format!(
        "Failed to connect to IAM server: {}",
        last_error.unwrap_or_else(|| "unknown connection error".to_string())
    ))
}

fn auth_connect_retry_delay(attempt: usize) -> Duration {
    AUTH_CONNECT_INITIAL_BACKOFF
        .saturating_mul(1u32 << attempt.min(4))
        .min(AUTH_CONNECT_MAX_BACKOFF)
}

fn retryable_connect_error(message: &str) -> bool {
    let lower = message.to_ascii_lowercase();
    [
        "transport error",
        "connection refused",
        "connection was not ready",
        "operation timed out",
        "deadline has elapsed",
        "dns error",
        "broken pipe",
        "connection reset",
    ]
    .iter()
    .any(|needle| lower.contains(needle))
}

fn prune_expired<T>(cache: &mut HashMap<String, CacheEntry<T>>) {
    let now = Instant::now();
    cache.retain(|_, entry| entry.expires_at > now);
}

fn cache_ttl_from_env(name: &str, default_ms: u64) -> Duration {
    std::env::var(name)
        .ok()
        .and_then(|value| value.parse::<u64>().ok())
        .map(Duration::from_millis)
        .unwrap_or_else(|| Duration::from_millis(default_ms))
}

fn authz_cache_key(
    principal: &Principal,
    action: &str,
    resource: &Resource,
) -> Result<String, serde_json::Error> {
    Ok(format!(
        "{}|{}|{}",
        action,
        serde_json::to_string(principal)?,
        serde_json::to_string(resource)?
    ))
}

fn resolve_tls_override(iam_endpoint: &str) -> Option<bool> {
    if iam_endpoint.starts_with("https://") {
        return Some(true);
    }
    if iam_endpoint.starts_with("http://") {
        return Some(false);
    }

    let env_val = std::env::var("PHOTON_IAM_TLS")
        .or_else(|_| std::env::var("IAM_CLIENT_TLS"))
        .ok()?;
    parse_env_bool(&env_val)
}

fn parse_env_bool(value: &str) -> Option<bool> {
    match value.trim().to_lowercase().as_str() {
        "1" | "true" | "yes" | "y" | "on" => Some(true),
        "0" | "false" | "no" | "n" | "off" => Some(false),
        _ => None,
    }
}

/// Extract tenant context from request extensions.
pub fn get_tenant_context<T>(request: &Request<T>) -> Result<TenantContext, Status> {
    request
        .extensions()
        .get::<TenantContext>()
        .cloned()
        .ok_or_else(|| Status::internal("Tenant context not found in request extensions"))
}

/// Resolve tenant IDs from the request, enforcing token match when provided.
pub fn resolve_tenant_ids<T>(
    request: &Request<T>,
    req_org_id: &str,
    req_project_id: &str,
) -> Result<(String, String), Status> {
    let tenant = get_tenant_context(request)?;
    resolve_tenant_ids_from_context(&tenant, req_org_id, req_project_id)
}

/// Resolve tenant IDs from an already-extracted tenant context.
pub fn resolve_tenant_ids_from_context(
    tenant: &TenantContext,
    req_org_id: &str,
    req_project_id: &str,
) -> Result<(String, String), Status> {
    if !req_org_id.is_empty() && req_org_id != tenant.org_id {
        return Err(Status::permission_denied("org_id does not match token"));
    }
    if !req_project_id.is_empty() && req_project_id != tenant.project_id {
        return Err(Status::permission_denied("project_id does not match token"));
    }

    Ok((tenant.org_id.clone(), tenant.project_id.clone()))
}

/// Build a resource scoped to a tenant.
pub fn resource_for_tenant(
    kind: impl Into<String>,
    id: impl Into<String>,
    org_id: impl Into<String>,
    project_id: impl Into<String>,
) -> Resource {
    Resource::new(kind, id, org_id, project_id)
}

fn extract_token_from_metadata(metadata: &MetadataMap) -> Result<String, Status> {
    let mut auth_err: Option<Status> = None;

    if let Some(auth_header) = metadata.get("authorization") {
        let auth_str = auth_header
            .to_str()
            .map_err(|_| Status::unauthenticated("Invalid authorization header encoding"))?;

        match parse_bearer_token(auth_str) {
            Ok(Some(token)) => return Ok(token),
            Ok(None) => auth_err = Some(Status::unauthenticated("Empty bearer token")),
            Err(err) => auth_err = Some(err),
        }
    }

    if let Some(photon_token_header) = metadata.get(PHOTON_AUTH_TOKEN_HEADER) {
        let token_str = photon_token_header
            .to_str()
            .map_err(|_| Status::unauthenticated("Invalid x-photon-auth-token header encoding"))?;

        if let Some(token) = parse_raw_token(token_str) {
            return Ok(token);
        }
    }

    if let Some(err) = auth_err {
        return Err(err);
    }

    Err(Status::unauthenticated(
        "Missing authorization or x-photon-auth-token header",
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn retryable_connect_error_matches_transport_failures() {
        assert!(retryable_connect_error("Internal error: transport error"));
        assert!(retryable_connect_error("connection was not ready"));
        assert!(retryable_connect_error("deadline has elapsed"));
        assert!(!retryable_connect_error("permission denied"));
    }

    #[test]
    fn auth_connect_retry_delay_is_capped() {
        assert_eq!(auth_connect_retry_delay(0), Duration::from_millis(500));
        assert_eq!(auth_connect_retry_delay(1), Duration::from_millis(1000));
        assert_eq!(auth_connect_retry_delay(2), Duration::from_millis(2000));
        assert_eq!(auth_connect_retry_delay(3), Duration::from_millis(4000));
        assert_eq!(auth_connect_retry_delay(4), Duration::from_secs(5));
        assert_eq!(auth_connect_retry_delay(8), Duration::from_secs(5));
    }
}

fn extract_token_from_headers(headers: &HeaderMap) -> Result<String, Status> {
    if let Some(auth_header) = headers.get(AUTHORIZATION) {
        let auth_str = auth_header
            .to_str()
            .map_err(|_| Status::unauthenticated("Invalid authorization header encoding"))?;

        if let Ok(Some(token)) = parse_bearer_token(auth_str) {
            return Ok(token);
        }

        let trimmed = auth_str.trim();
        if !trimmed.is_empty() && trimmed.split_whitespace().count() == 1 {
            return Ok(trimmed.to_string());
        }
    }

    if let Some(photon_header) = headers.get(PHOTON_AUTH_TOKEN_HEADER) {
        let token_str = photon_header
            .to_str()
            .map_err(|_| Status::unauthenticated("Invalid x-photon-auth-token header encoding"))?;

        if let Some(token) = parse_raw_token(token_str) {
            return Ok(token);
        }
    }

    Err(Status::unauthenticated(
        "Missing authorization or x-photon-auth-token header",
    ))
}

fn parse_bearer_token(value: &str) -> Result<Option<String>, Status> {
    let mut parts = value.trim().split_whitespace();
    let scheme = parts
        .next()
        .ok_or_else(|| Status::unauthenticated("Empty bearer token"))?;

    if !scheme.eq_ignore_ascii_case("bearer") {
        return Err(Status::unauthenticated(
            "Authorization header must use Bearer scheme",
        ));
    }

    let token = parts.next().unwrap_or("").trim();
    if token.is_empty() {
        return Ok(None);
    }

    if parts.next().is_some() {
        return Err(Status::unauthenticated(
            "Authorization header must be in the form: Bearer <token>",
        ));
    }

    Ok(Some(token.to_string()))
}

fn parse_raw_token(value: &str) -> Option<String> {
    let trimmed = value.trim();
    if trimmed.is_empty() {
        return None;
    }

    if let Ok(Some(token)) = parse_bearer_token(trimmed) {
        return Some(token);
    }

    Some(trimmed.to_string())
}
