use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use p256::ecdsa::SigningKey;
use rand::rngs::OsRng;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
use thiserror::Error;
use tranquil_oauth::{
    AuthorizationServerMetadata, ClientMetadata, compute_es256_jkt, compute_pkce_challenge,
    create_dpop_proof,
};
use tranquil_types::{
    AuthorizationCode, ClientId, CrossPdsState, Did, Issuer, PdsUrl, ReachPolicy,
};

use crate::cache::Cache;

const SERVER_METADATA_TTL: Duration = Duration::from_secs(300);

#[derive(Error, Debug)]
pub enum CrossPdsError {
    #[error("failed to fetch OAuth metadata: {0}")]
    MetadataFetch(String),
    #[error("controller PDS has no PAR endpoint")]
    NoParEndpoint,
    #[error("PAR request failed: {0}")]
    ParFailed(String),
    #[error("token exchange failed: {0}")]
    TokenExchangeFailed(String),
    #[error("invalid token response: {0}")]
    InvalidTokenResponse(String),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CrossPdsAuthState {
    pub original_request_uri: String,
    pub controller_did: Did,
    pub controller_pds_url: PdsUrl,
    pub code_verifier: String,
    pub dpop_private_key_der: String,
    pub delegated_did: Did,
    pub expected_issuer: Option<Issuer>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ParResult {
    pub request_uri: String,
    pub authorize_url: String,
}

pub struct DelegationOAuthUrls {
    pub client_id: ClientId,
    pub redirect_uri: String,
}

pub fn delegation_oauth_urls(hostname: &str) -> DelegationOAuthUrls {
    DelegationOAuthUrls {
        client_id: ClientId::new(format!(
            "https://{}/oauth/delegation/client-metadata",
            hostname
        )),
        redirect_uri: format!("https://{}/oauth/delegation/callback", hostname),
    }
}

pub struct CrossPdsOAuthClient {
    http: Client,
    cache: Arc<dyn Cache>,
}

impl CrossPdsOAuthClient {
    pub fn new(cache: Arc<dyn Cache>, fetch_policy: ReachPolicy) -> Self {
        let http = Client::builder()
            .timeout(Duration::from_secs(15))
            .connect_timeout(Duration::from_secs(5))
            .redirect(tranquil_types::redirect_policy(fetch_policy))
            .dns_resolver(tranquil_types::dns_guard(fetch_policy))
            .build()
            .expect("failed to build cross-PDS OAuth HTTP client");
        Self { http, cache }
    }

    pub async fn store_auth_state(
        &self,
        state_key: &CrossPdsState,
        auth_state: &CrossPdsAuthState,
    ) -> Result<(), CrossPdsError> {
        let cache_key = crate::cache_keys::cross_pds_state_key(state_key);
        let json_bytes = serde_json::to_vec(auth_state)
            .map_err(|e| CrossPdsError::ParFailed(format!("serialize auth state: {}", e)))?;
        let encrypted = crate::config::encrypt_key(&json_bytes)
            .map_err(|e| CrossPdsError::ParFailed(format!("encrypt auth state: {}", e)))?;
        self.cache
            .set_bytes(&cache_key, &encrypted, Duration::from_secs(600))
            .await
            .map_err(|e| CrossPdsError::ParFailed(format!("cache auth state: {}", e)))
    }

    pub async fn retrieve_auth_state(
        &self,
        state_key: &CrossPdsState,
    ) -> Result<CrossPdsAuthState, CrossPdsError> {
        let cache_key = crate::cache_keys::cross_pds_state_key(state_key);
        let encrypted_bytes = self.cache.get_bytes(&cache_key).await.ok_or_else(|| {
            CrossPdsError::TokenExchangeFailed("auth state expired or not found".into())
        })?;
        let _ = self.cache.delete(&cache_key).await;
        let decrypted =
            crate::config::decrypt_key(&encrypted_bytes, Some(crate::config::ENCRYPTION_VERSION))
                .map_err(|e| {
                CrossPdsError::TokenExchangeFailed(format!("decrypt auth state: {}", e))
            })?;
        serde_json::from_slice(&decrypted).map_err(|e| {
            CrossPdsError::TokenExchangeFailed(format!("deserialize auth state: {}", e))
        })
    }

    pub async fn check_remote_is_delegated(&self, pds_url: &PdsUrl, did: &Did) -> Option<bool> {
        let mut url = pds_url.endpoint("oauth/security-status");
        url.query_pairs_mut()
            .append_pair("identifier", did.as_str());
        let resp = self.http.get(url).send().await.ok()?;
        if !resp.status().is_success() {
            return None;
        }
        #[derive(Deserialize)]
        #[serde(rename_all = "camelCase")]
        struct RemoteSecurityStatus {
            is_delegated: Option<bool>,
        }
        resp.json::<RemoteSecurityStatus>()
            .await
            .ok()
            .and_then(|s| s.is_delegated)
    }

    async fn send_with_dpop_retry(
        &self,
        signing_key: &SigningKey,
        method: &str,
        url: &str,
        params: &[(&str, String)],
        access_token_hash: Option<&str>,
    ) -> Result<reqwest::Response, String> {
        let make_proof = |nonce: Option<&str>| {
            create_dpop_proof(signing_key, method, url, nonce, access_token_hash)
                .map_err(|e| format!("{:?}", e))
        };

        let resp = self
            .http
            .post(url)
            .header("DPoP", &make_proof(None)?)
            .form(params)
            .send()
            .await
            .map_err(|e| e.to_string())?;

        let nonce = resp
            .headers()
            .get("dpop-nonce")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string());
        let needs_retry = matches!(
            resp.status(),
            reqwest::StatusCode::BAD_REQUEST | reqwest::StatusCode::UNAUTHORIZED
        );

        if needs_retry && nonce.is_some() {
            return self
                .http
                .post(url)
                .header("DPoP", &make_proof(nonce.as_deref())?)
                .form(params)
                .send()
                .await
                .map_err(|e| e.to_string());
        }
        Ok(resp)
    }

    async fn resolve_authorization_server(
        &self,
        pds_url: &PdsUrl,
    ) -> Result<Issuer, CrossPdsError> {
        let resource_url = pds_url.endpoint(".well-known/oauth-protected-resource");
        if let Ok(resp) = self.http.get(resource_url).send().await
            && resp.status().is_success()
        {
            #[derive(Deserialize)]
            struct ProtectedResource {
                authorization_servers: Option<Vec<String>>,
            }
            if let Ok(pr) = resp.json::<ProtectedResource>().await
                && let Some(server) = pr.authorization_servers.and_then(|s| s.into_iter().next())
            {
                return Issuer::new(server)
                    .map_err(|e| CrossPdsError::MetadataFetch(e.to_string()));
            }
        }
        Issuer::new(pds_url.as_str()).map_err(|e| CrossPdsError::MetadataFetch(e.to_string()))
    }

    pub async fn fetch_server_metadata(
        &self,
        pds_url: &PdsUrl,
    ) -> Result<AuthorizationServerMetadata, CrossPdsError> {
        crate::cache::cached_json(
            self.cache.as_ref(),
            &crate::cache_keys::cross_pds_oauth_meta_key(pds_url),
            SERVER_METADATA_TTL,
            || self.fetch_verified_server_metadata(pds_url),
        )
        .await
    }

    async fn fetch_verified_server_metadata(
        &self,
        pds_url: &PdsUrl,
    ) -> Result<AuthorizationServerMetadata, CrossPdsError> {
        let auth_server = self.resolve_authorization_server(pds_url).await?;

        let url = auth_server.endpoint(".well-known/oauth-authorization-server");
        let resp = self
            .http
            .get(url.clone())
            .send()
            .await
            .map_err(|e| CrossPdsError::MetadataFetch(e.to_string()))?;

        if !resp.status().is_success() {
            return Err(CrossPdsError::MetadataFetch(format!(
                "HTTP {} from {}",
                resp.status(),
                url
            )));
        }

        let meta: AuthorizationServerMetadata = resp
            .json()
            .await
            .map_err(|e| CrossPdsError::MetadataFetch(e.to_string()))?;

        if meta.issuer != auth_server {
            return Err(CrossPdsError::MetadataFetch(format!(
                "issuer mismatch: {} serves metadata for {}",
                auth_server, meta.issuer
            )));
        }

        Ok(meta)
    }

    pub async fn initiate_par(
        &self,
        pds_url: &PdsUrl,
        urls: &DelegationOAuthUrls,
        login_hint: Option<&str>,
        original_request_uri: &str,
        controller_did: &Did,
        delegated_did: &Did,
    ) -> Result<(ParResult, CrossPdsAuthState, CrossPdsState), CrossPdsError> {
        let meta = self.fetch_server_metadata(pds_url).await?;
        let par_endpoint = meta
            .pushed_authorization_request_endpoint
            .as_ref()
            .ok_or(CrossPdsError::NoParEndpoint)?;

        let code_verifier = crate::util::generate_random_token();
        let code_challenge = compute_pkce_challenge(&code_verifier);
        let state = CrossPdsState::new(crate::util::generate_random_token());

        let signing_key = SigningKey::random(&mut OsRng);
        let dpop_key_der = URL_SAFE_NO_PAD.encode(signing_key.to_bytes());

        let dpop_jkt = compute_es256_jkt(&signing_key)
            .map_err(|e| CrossPdsError::ParFailed(format!("{:?}", e)))?;

        let mut params = vec![
            ("response_type", "code".to_string()),
            ("client_id", urls.client_id.to_string()),
            ("redirect_uri", urls.redirect_uri.clone()),
            ("scope", "atproto".to_string()),
            ("state", state.to_string()),
            ("code_challenge", code_challenge),
            ("code_challenge_method", "S256".to_string()),
            ("dpop_jkt", dpop_jkt),
        ];
        if let Some(hint) = login_hint {
            params.push(("login_hint", hint.to_string()));
        }

        let resp = self
            .send_with_dpop_retry(&signing_key, "POST", par_endpoint.as_str(), &params, None)
            .await
            .map_err(|e| CrossPdsError::ParFailed(e.to_string()))?;

        if !resp.status().is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(CrossPdsError::ParFailed(format!("PAR rejected: {}", body)));
        }

        #[derive(Deserialize)]
        struct ParResp {
            request_uri: String,
        }

        let par_resp: ParResp = resp
            .json()
            .await
            .map_err(|e| CrossPdsError::ParFailed(e.to_string()))?;

        let mut authorize_url = meta.authorization_endpoint.url().clone();
        authorize_url
            .query_pairs_mut()
            .append_pair("request_uri", &par_resp.request_uri)
            .append_pair("client_id", &urls.client_id);

        let auth_state = CrossPdsAuthState {
            original_request_uri: original_request_uri.to_string(),
            controller_did: controller_did.clone(),
            controller_pds_url: pds_url.clone(),
            code_verifier,
            dpop_private_key_der: dpop_key_der,
            delegated_did: delegated_did.clone(),
            expected_issuer: Some(meta.issuer.clone()),
        };

        Ok((
            ParResult {
                request_uri: par_resp.request_uri,
                authorize_url: authorize_url.into(),
            },
            auth_state,
            state,
        ))
    }

    pub async fn exchange_code(
        &self,
        auth_state: &CrossPdsAuthState,
        code: &AuthorizationCode,
        client_id: &ClientId,
        redirect_uri: &str,
    ) -> Result<String, CrossPdsError> {
        let meta = self
            .fetch_server_metadata(&auth_state.controller_pds_url)
            .await?;

        let key_bytes = URL_SAFE_NO_PAD
            .decode(&auth_state.dpop_private_key_der)
            .map_err(|e| CrossPdsError::TokenExchangeFailed(e.to_string()))?;
        let signing_key = SigningKey::from_bytes((&key_bytes[..]).into())
            .map_err(|e| CrossPdsError::TokenExchangeFailed(e.to_string()))?;

        let params = vec![
            ("grant_type", "authorization_code".to_string()),
            ("code", code.to_string()),
            ("redirect_uri", redirect_uri.to_string()),
            ("code_verifier", auth_state.code_verifier.clone()),
            ("client_id", client_id.to_string()),
        ];

        let resp = self
            .send_with_dpop_retry(
                &signing_key,
                "POST",
                meta.token_endpoint.as_str(),
                &params,
                None,
            )
            .await
            .map_err(CrossPdsError::TokenExchangeFailed)?;

        if !resp.status().is_success() {
            let body = resp.text().await.unwrap_or_default();
            return Err(CrossPdsError::TokenExchangeFailed(format!(
                "Token exchange rejected: {}",
                body
            )));
        }

        #[derive(Deserialize)]
        struct TokenResp {
            sub: Option<String>,
            token_type: Option<String>,
            error: Option<String>,
            error_description: Option<String>,
        }

        let token_resp: TokenResp = resp
            .json()
            .await
            .map_err(|e| CrossPdsError::InvalidTokenResponse(e.to_string()))?;

        if let Some(ref err) = token_resp.error {
            let desc = token_resp.error_description.as_deref().unwrap_or("unknown");
            return Err(CrossPdsError::TokenExchangeFailed(format!(
                "{}: {}",
                err, desc
            )));
        }

        if let Some(ref tt) = token_resp.token_type
            && !tt.eq_ignore_ascii_case("DPoP")
        {
            return Err(CrossPdsError::InvalidTokenResponse(format!(
                "expected token_type DPoP, got {}",
                tt
            )));
        }

        token_resp
            .sub
            .ok_or_else(|| CrossPdsError::InvalidTokenResponse("missing sub claim".to_string()))
    }
}

pub fn build_client_metadata(hostname: &str) -> ClientMetadata {
    let urls = delegation_oauth_urls(hostname);
    ClientMetadata {
        client_id: urls.client_id,
        client_name: Some(hostname.to_string()),
        client_uri: Some(format!("https://{}", hostname)),
        redirect_uris: vec![urls.redirect_uri],
        grant_types: vec!["authorization_code".to_string()],
        response_types: vec!["code".to_string()],
        scope: Some("atproto".to_string()),
        dpop_bound_access_tokens: Some(true),
        token_endpoint_auth_method: Some("none".to_string()),
        application_type: Some("web".to_string()),
        ..ClientMetadata::default()
    }
}
