//! HTTPS-only fetcher for `/api/plugin/public-keys` that verifies the
//! JWS envelope against the bundled root anchors and caches keys per
//! tenant+kid. Mirrors `pay-plugin-core/src/publicKeyResolver.ts`.

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

use ed25519_dalek::pkcs8::DecodePublicKey;
use ed25519_dalek::VerifyingKey;
use reqwest::blocking::Client;
use serde::Deserialize;

use crate::jws::verify_ed25519_pem;
use crate::root_keys::root_key_pem;

#[derive(Debug)]
pub enum PublicKeyResolverError {
    InsecureBaseUrl(String),
    InvalidBaseUrl(String),
    Http(String),
    UnexpectedStatus(u16),
    InvalidPayload(String),
    UnknownRootKid(String),
    JwsRequired,
    JwsInvalid(String),
    KeyNotFound { tenant_id: String, kid: Option<String> },
}

#[derive(Debug, Deserialize)]
struct EnvelopeKey {
    kid: String,
    algorithm: String,
    #[allow(dead_code)]
    status: String,
    public_key_pem: String,
}

#[derive(Debug, Deserialize)]
struct PlainEnvelope {
    tenant_id: String,
    keys: Vec<EnvelopeKey>,
}

#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum WireResponse {
    Jws {
        format: String,
        root_kid: String,
        compact: String,
    },
    Plain(PlainEnvelope),
}

struct CacheEntry {
    verifying_key: VerifyingKey,
    fetched_at: Instant,
}

pub struct PublicKeyResolverConfig {
    pub base_url: String,
    pub allow_insecure: bool,
    pub refresh_interval: Duration,
    /// kid -> SPKI PEM. Defaults to the bundled anchors when empty.
    pub root_public_keys: HashMap<String, String>,
}

impl Default for PublicKeyResolverConfig {
    fn default() -> Self {
        Self {
            base_url: String::new(),
            allow_insecure: false,
            refresh_interval: Duration::from_secs(5 * 60),
            root_public_keys: HashMap::new(),
        }
    }
}

pub struct PublicKeyResolver {
    config: PublicKeyResolverConfig,
    client: Client,
    cache: Mutex<HashMap<String, CacheEntry>>,
}

impl PublicKeyResolver {
    pub fn new(config: PublicKeyResolverConfig) -> Result<Self, PublicKeyResolverError> {
        let parsed = url::Url::parse(&config.base_url)
            .map_err(|e| PublicKeyResolverError::InvalidBaseUrl(e.to_string()))?;
        match parsed.scheme() {
            "https" => {}
            "http" if config.allow_insecure => {}
            other => {
                return Err(PublicKeyResolverError::InsecureBaseUrl(other.to_string()));
            }
        }

        let client = Client::builder()
            .https_only(!config.allow_insecure)
            .build()
            .map_err(|e| PublicKeyResolverError::Http(e.to_string()))?;

        Ok(Self {
            config,
            client,
            cache: Mutex::new(HashMap::new()),
        })
    }

    pub fn verifying_key_for(
        &self,
        tenant_id: &str,
        kid: &str,
    ) -> Result<VerifyingKey, PublicKeyResolverError> {
        let cache_key = format!("{tenant_id}::{kid}");
        if let Some(entry) = self.cache.lock().unwrap().get(&cache_key) {
            if entry.fetched_at.elapsed() < self.config.refresh_interval {
                return Ok(entry.verifying_key);
            }
        }

        self.refresh(QueryParam::Tenant(tenant_id.to_string()))?;

        let cache = self.cache.lock().unwrap();
        cache
            .get(&cache_key)
            .map(|entry| entry.verifying_key)
            .ok_or_else(|| PublicKeyResolverError::KeyNotFound {
                tenant_id: tenant_id.to_string(),
                kid: Some(kid.to_string()),
            })
    }

    pub fn verifying_key_for_product(
        &self,
        product_id: &str,
    ) -> Result<(String, VerifyingKey), PublicKeyResolverError> {
        let tenant_id = self.refresh(QueryParam::Product(product_id.to_string()))?;
        let prefix = format!("{tenant_id}::");
        let cache = self.cache.lock().unwrap();
        for (cache_key, entry) in cache.iter() {
            if cache_key.starts_with(&prefix) {
                return Ok((tenant_id.clone(), entry.verifying_key));
            }
        }
        Err(PublicKeyResolverError::KeyNotFound {
            tenant_id,
            kid: None,
        })
    }

    fn refresh(&self, query: QueryParam) -> Result<String, PublicKeyResolverError> {
        let mut url = url::Url::parse(&self.config.base_url)
            .map_err(|e| PublicKeyResolverError::InvalidBaseUrl(e.to_string()))?;
        url.set_path("/api/plugin/public-keys");
        {
            let mut pairs = url.query_pairs_mut();
            match &query {
                QueryParam::Tenant(value) => {
                    pairs.append_pair("tenant_id", value);
                }
                QueryParam::Product(value) => {
                    pairs.append_pair("product_id", value);
                }
            }
        }

        let response = self
            .client
            .get(url)
            .send()
            .map_err(|e| PublicKeyResolverError::Http(e.to_string()))?;
        let status = response.status();
        if !status.is_success() {
            return Err(PublicKeyResolverError::UnexpectedStatus(status.as_u16()));
        }
        let wire: WireResponse = response
            .json()
            .map_err(|e| PublicKeyResolverError::InvalidPayload(e.to_string()))?;

        let plain = match wire {
            WireResponse::Jws {
                format,
                root_kid,
                compact,
            } => {
                if format != "jws" {
                    return Err(PublicKeyResolverError::InvalidPayload(format!(
                        "unexpected format {format}"
                    )));
                }
                let root_pem = self.resolve_root_pem(&root_kid)?;
                let verified = verify_ed25519_pem(&compact, &root_pem)
                    .map_err(|e| PublicKeyResolverError::JwsInvalid(format!("{e:?}")))?;
                serde_json::from_slice::<PlainEnvelope>(&verified.payload)
                    .map_err(|e| PublicKeyResolverError::InvalidPayload(e.to_string()))?
            }
            WireResponse::Plain(plain) => {
                if !self.config.root_public_keys.is_empty() || !root_keys_empty() {
                    return Err(PublicKeyResolverError::JwsRequired);
                }
                plain
            }
        };

        let now = Instant::now();
        let mut cache = self.cache.lock().unwrap();
        for key in plain.keys {
            if key.algorithm != "EdDSA" || key.public_key_pem.is_empty() {
                continue;
            }
            let verifying_key = match VerifyingKey::from_public_key_pem(&key.public_key_pem) {
                Ok(value) => value,
                Err(_) => continue,
            };
            cache.insert(
                format!("{}::{}", plain.tenant_id, key.kid),
                CacheEntry {
                    verifying_key,
                    fetched_at: now,
                },
            );
        }
        Ok(plain.tenant_id)
    }

    fn resolve_root_pem(&self, kid: &str) -> Result<String, PublicKeyResolverError> {
        if let Some(pem) = self.config.root_public_keys.get(kid) {
            return Ok(pem.clone());
        }
        root_key_pem(kid)
            .map(str::to_owned)
            .ok_or_else(|| PublicKeyResolverError::UnknownRootKid(kid.to_string()))
    }
}

fn root_keys_empty() -> bool {
    crate::root_keys::AXIS_PLATFORM_ROOT_PUBLIC_KEYS.is_empty()
}

enum QueryParam {
    Tenant(String),
    Product(String),
}
