use std::collections::HashMap;

use futures::future::BoxFuture;
use oauth2::{RefreshToken, TokenResponse as _};
use rmcp::transport::auth::{
    AuthClient, AuthorizationManager, CredentialStore, InMemoryCredentialStore, OAuthClientConfig,
    OAuthState, StoredCredentials,
};
use rmcp::transport::{AuthError, AuthorizationSession};
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use url::Url;
use uuid::Uuid;
use warp_core::channel::ChannelState;
use warpui_extras::secure_storage::AppContextExt as _;

pub const TEMPLATABLE_MCP_CREDENTIALS_KEY: &str = "TemplatableMcpCredentials";
pub const FILE_BASED_MCP_CREDENTIALS_KEY: &str = "FileBasedMcpCredentials";

/// The issuer URL for GitHub's OAuth provider.
const GITHUB_ISSUER: &str = "https://github.com/login/oauth";

static GITHUB_OAUTH_SCOPES: [&str; 7] = [
    "repo",
    "read:org",
    "gist",
    "notifications",
    "user",
    "project",
    "workflow",
];

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PersistedCredentials {
    /// The credential information that `rmcp` wants us to store and retrieve.
    #[serde(flatten)]
    credentials: StoredCredentials,
    /// The client secret for the OAuth application.
    ///
    /// This is needed to properly refresh tokens when using DCR (Dynamic Client Registration),
    /// as the server expects the client to provide the secret when refreshing.
    client_secret: Option<String>,
}

/// Maps cloud MCP installation UUID to its OAuth credentials in secure storage.
pub type PersistedCredentialsMap = HashMap<Uuid, PersistedCredentials>;

// Maps a consistent hash of the installation to its persisted credentials
pub type FileBasedPersistedCredentialsMap = HashMap<u64, PersistedCredentials>;
pub type PersistCredentialsCallback =
    Box<dyn Fn(Uuid, PersistedCredentials) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
pub type RequiresAuthenticationCallback =
    Box<dyn Fn(Uuid, String, String) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
pub type AuthenticatedCallback =
    Box<dyn Fn(String) -> BoxFuture<'static, anyhow::Result<()>> + Send>;

/// A credential store that wraps [`InMemoryCredentialStore`] and persists token
/// updates to Warp's secure storage via a channel.
///
/// When rmcp auto-refreshes an expired access token at runtime, the rotated
/// tokens are only saved to the in-memory store by default. This wrapper
/// ensures they also get written back to secure storage so they survive app
/// restarts.
struct PersistingCredentialStore {
    inner: InMemoryCredentialStore,
    client_secret: Option<String>,
    persist_tx: async_channel::Sender<PersistedCredentials>,
}

impl PersistingCredentialStore {
    /// Per RFC 6749 §6, the authorization server MAY issue a new refresh token on
    /// refresh, but is not required to. Many OAuth providers (e.g. Figma) only
    /// issue a refresh token on the initial authorization grant and omit it from
    /// subsequent refresh responses. If we blindly persist the new token response,
    /// the refresh token is lost and the next session (or next in-process refresh)
    /// requires a full re-auth.
    ///
    /// When the new response omits a refresh token, carry forward the one already
    /// in the store. See: <https://datatracker.ietf.org/doc/html/rfc6749#section-6>
    async fn apply_refresh_token_carry_forward(&self, credentials: &mut StoredCredentials) {
        if credentials
            .token_response
            .as_ref()
            .is_none_or(|tr| tr.refresh_token().is_some())
        {
            return;
        }

        if let Some(prev_rt) = self
            .inner
            .load()
            .await
            .ok()
            .and_then(|opt| opt)
            .and_then(|prev| prev.token_response)
            .and_then(|prev_tr| prev_tr.refresh_token().cloned())
        {
            if let Some(tr) = credentials.token_response.as_mut() {
                // Carry forward the existing/previous refresh token, constructing new if needed
                tr.set_refresh_token(Some(RefreshToken::new(prev_rt.secret().to_string())));
            }
        }
    }
}

#[async_trait::async_trait]
impl CredentialStore for PersistingCredentialStore {
    async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
        self.inner.load().await
    }

    async fn save(&self, mut credentials: StoredCredentials) -> Result<(), AuthError> {
        self.apply_refresh_token_carry_forward(&mut credentials)
            .await;

        self.inner.save(credentials.clone()).await?;

        // Only persist credentials if we actually have any.
        if credentials.token_response.is_some() {
            let _ = self.persist_tx.try_send(PersistedCredentials {
                credentials,
                client_secret: self.client_secret.clone(),
            });
        }
        Ok(())
    }

    async fn clear(&self) -> Result<(), AuthError> {
        self.inner.clear().await
    }
}

/// Installs a [`PersistingCredentialStore`] on the given auth manager so that
/// runtime token auto-refreshes are written back to Warp's secure storage.
///
/// A background tokio task is spawned to receive credential updates and persist
/// them via the provided callback. The task terminates when the auth manager (and
/// thus the credential store's sender) is dropped.
///
/// Note: this store is not responsible for the initial population of credentials.
/// Instead, the caller seeds the inner store with any existing credentials prior
/// to installation (see [`install_persisting_credential_store`]). This store's
/// sole role is to write token updates back to secure storage as they occur.
async fn install_persisting_credential_store(
    auth_manager: &mut AuthorizationManager,
    persisted_credentials: Option<PersistedCredentials>,
    persist_credentials: PersistCredentialsCallback,
    installation_uuid: Uuid,
) {
    let client_secret = persisted_credentials
        .as_ref()
        .and_then(|c| c.client_secret.clone());
    let in_memory_store = InMemoryCredentialStore::new();

    // If we have persisted credentials, populate the backing in-memory store with them.
    if let Some(credentials) = persisted_credentials {
        let _ = in_memory_store.save(credentials.credentials).await;
    }

    let (persist_tx, persist_rx) = async_channel::unbounded();
    let store = PersistingCredentialStore {
        inner: in_memory_store,
        client_secret,
        persist_tx,
    };

    auth_manager.set_credential_store(store);

    tokio::spawn(async move {
        while let Ok(credentials) = persist_rx.recv().await {
            if let Err(err) = persist_credentials(installation_uuid, credentials).await {
                log::warn!("Failed to persist auto-refreshed MCP credentials: {err:?}");
            }
        }
    });
}

/// Context for OAuth authentication flows.
pub struct AuthContext {
    pub oauth_result_rx: async_channel::Receiver<CallbackResult>,
    pub uuid: Uuid,
    pub persisted_credentials: Option<PersistedCredentials>,
    /// Whether the client is running in headless/CLI mode.
    pub is_headless: bool,
    /// Whether this server was auto-discovered from a repo MCP configuration file.
    pub is_file_based: bool,
    pub persist_credentials: PersistCredentialsCallback,
    pub requires_authentication: RequiresAuthenticationCallback,
    pub authenticated: Option<AuthenticatedCallback>,
}

/// Result of OAuth callback.
#[derive(Debug, Clone)]
pub enum CallbackResult {
    Success { code: String, csrf_token: String },
    Error { error: Option<String> },
}

/// Makes an authenticated client for the given authorization server.
///
/// This takes in the URL of the resource to authenticate for, and uses that
/// to determine the authorization server.
///
/// Upon success, returns the client and a boolean indicating whether the user was required to
/// re-authenticate (e.g. re-log in).
pub async fn make_authenticated_client(
    resource_url: &str,
    auth_context: AuthContext,
) -> Result<(AuthClient<reqwest::Client>, bool), AuthError> {
    let AuthContext {
        oauth_result_rx,
        uuid,
        persisted_credentials,
        is_headless,
        is_file_based,
        persist_credentials,
        requires_authentication,
        ..
    } = auth_context;

    // Build the redirect URI using the channel's URL scheme.
    // Routing data (the server UUID) is passed via the OAuth `state` parameter instead
    // of the redirect URI so that the URI exactly matches what is registered during
    // Dynamic Client Registration, satisfying RFC 6749 §3.1.2.2 exact-match validation.
    let redirect_uri = format!("{}://mcp/oauth2callback", ChannelState::url_scheme());

    // Create the auth manager and initialize it with a backing credential store that persists
    // new credentials to secure storage.
    let client_id = persisted_credentials
        .as_ref()
        .map(|c| c.credentials.client_id.clone());
    let client_secret = persisted_credentials
        .as_ref()
        .and_then(|c| c.client_secret.clone());
    let mut auth_manager = AuthorizationManager::new(resource_url).await?;
    install_persisting_credential_store(
        &mut auth_manager,
        persisted_credentials,
        persist_credentials,
        uuid,
    )
    .await;

    // If we loaded persisted credentials from the store, and we have a valid access token
    // (or successfully refreshed a valid refresh token), we're already authorized and good
    // to go.
    if auth_manager.initialize_from_store().await? && auth_manager.get_access_token().await.is_ok()
    {
        if let (Some(client_id), Some(client_secret)) = (client_id, client_secret) {
            auth_manager.configure_client(
                OAuthClientConfig::new(client_id, redirect_uri.clone())
                    .with_client_secret(client_secret),
            )?;
        }
        return Ok((AuthClient::new(reqwest::Client::new(), auth_manager), false));
    }

    // If we're in headless mode and we reach here, it means we either have no credentials
    // or the cached credentials failed to refresh. Block interactive OAuth in headless mode.
    if is_headless {
        if is_file_based {
            log::warn!(
                "File-based MCP server {uuid} requires OAuth authentication; \
                 skipping in headless mode. To use this server, authenticate it \
                 in the Warp desktop app first."
            );
        }
        return Err(AuthError::AuthorizationFailed(
            "MCP server requires OAuth authentication. Please authenticate this server in the \
             Warp desktop app first, then try again."
                .to_string(),
        ));
    }

    let metadata = auth_manager.discover_metadata().await?;

    // Configure the auth manager's OAuth client using dynamic or static client registration.
    let mut oauth_state = if let Some(provider) = metadata
        .issuer
        .as_deref()
        .and_then(ChannelState::mcp_oauth_provider_by_issuer)
    {
        // Configure the auth manager based on the static MCP configuration for this
        // issuer.
        auth_manager.set_metadata(metadata);

        let scopes = if provider.issuer == GITHUB_ISSUER {
            GITHUB_OAUTH_SCOPES
                .into_iter()
                .map(ToString::to_string)
                .collect()
        } else {
            vec![]
        };
        auth_manager.configure_client(
            OAuthClientConfig::new(provider.client_id, redirect_uri.clone())
                .with_client_secret(provider.client_secret)
                .with_scopes(scopes),
        )?;

        // We do a scope "upgrade" with no additional scopes here as it's the easiest way
        // to construct an auth URL.
        let auth_url = auth_manager.request_scope_upgrade("").await?;
        OAuthState::Session(AuthorizationSession::for_scope_upgrade(
            auth_manager,
            auth_url,
            &redirect_uri,
        ))
    } else {
        // Try dynamic client registration.
        let mut oauth_state = OAuthState::Unauthorized(auth_manager);
        oauth_state
            .start_authorization(&[], &redirect_uri, Some("Warp"))
            .await?;
        oauth_state
    };

    let auth_url = oauth_state.get_authorization_url().await?;

    // Extract the CSRF token that rmcp embedded as the `state` query parameter in the
    // authorization URL. We register a csrf→uuid mapping on the manager so that
    // `handle_oauth_callback` can route the callback to the right server without
    // relying on `server_id` being present in the redirect URI.
    let csrf_state = Url::parse(&auth_url)
        .ok()
        .and_then(|u| {
            u.query_pairs()
                .find(|(k, _)| k == "state")
                .map(|(_, v)| v.into_owned())
        })
        .unwrap_or_default();

    if let Err(err) = requires_authentication(uuid, csrf_state, auth_url).await {
        log::warn!("Failed to emit RequiresAuthentication state: {err:?}");
    }

    // Wait for the authorization code from the OAuth callback channel.
    let oauth_result = oauth_result_rx
        .recv()
        .await
        .map_err(|e| AuthError::InternalError(e.to_string()))?;

    let (code, csrf_token) = match &oauth_result {
        CallbackResult::Success { code, csrf_token } => (code, csrf_token),
        CallbackResult::Error { error } => {
            return Err(AuthError::AuthorizationFailed(
                error.as_deref().unwrap_or("unknown error").to_string(),
            ));
        }
    };

    // Handle the callback with the received authorization code and CSRF token.
    oauth_state.handle_callback(code, csrf_token).await?;

    let auth_manager = oauth_state.into_authorization_manager().ok_or_else(|| {
        AuthError::InternalError("Failed to create authorization manager".to_string())
    })?;

    Ok((AuthClient::new(reqwest::Client::new(), auth_manager), true))
}

/// Loads credentials from secure storage at the provided key.
pub fn load_credentials_from_secure_storage<T: DeserializeOwned + Default>(
    app: &mut warpui::AppContext,
    key: &str,
) -> T {
    app.secure_storage()
        .read_value(key)
        .inspect_err(|err| {
            if !matches!(err, warpui_extras::secure_storage::Error::NotFound) {
                log::warn!("Failed to read MCP credentials from secure storage: {err:#}");
            }
        })
        .ok()
        .and_then(|value| serde_json::from_str(&value).ok())
        .unwrap_or_default()
}

/// Writes credentials to secure storage at the provided key.
pub fn write_to_secure_storage<T: Serialize>(
    app: &mut warpui::AppContext,
    key: &str,
    credentials: &T,
) {
    match serde_json::to_string(credentials) {
        Ok(json) => {
            app.secure_storage()
                .write_value(key, &json)
                .inspect_err(|err| {
                    log::error!("Failed to write MCP credentials to secure storage: {err:#}")
                })
                .ok();
        }
        Err(err) => {
            log::error!("Failed to serialize MCP credentials for secure storage: {err:#}");
        }
    }
}

#[cfg(test)]
#[path = "oauth_tests.rs"]
mod tests;