//! `GET /oidc/callback?code=&state=` — browser-facing redirect target.
//!
//! After the user signs in at the IdP, the IdP redirects their browser
//! here. We exchange the IdP code for tokens, fetch userinfo, find/create
//! a local user, mint our access token, and mark the session `success`.
//! The browser sees a small "you can close this window" page; the desktop
//! client picks up the token via `/api/oidc/auth-query`.

use crate::api::admin::oidc_login::ADMIN_SENTINEL;
use crate::api::auth::mint_token;
use crate::api::middleware::{sha256_token, SESSION_COOKIE};
use crate::api::oidc::{discovery, require_provider};
use crate::api::state::AppState;
use axum::extract::{Extension, Query};
use axum::http::header::{LOCATION, SET_COOKIE};
use axum::http::{HeaderMap, HeaderValue, StatusCode};
use axum::response::{Html, IntoResponse, Response};
use serde::Deserialize;
use serde_json::Value;
use std::sync::Arc;

#[derive(Debug, Deserialize)]
pub struct CallbackQuery {
    #[serde(default)]
    pub code: String,
    #[serde(default)]
    pub state: String,
    /// Some IdPs forward an error here on failed auth (e.g. user clicked
    /// "deny"). We surface it as the session error and as a friendly page.
    #[serde(default)]
    pub error: Option<String>,
    #[serde(default)]
    pub error_description: Option<String>,
}

pub async fn callback(
    Extension(state): Extension<Arc<AppState>>,
    Query(q): Query<CallbackQuery>,
) -> Response {
    match handle(state.clone(), q).await {
        Ok(ok) if ok.is_admin_flow => {
            if !ok.user_is_admin {
                return Html(html_page(
                    "Sign-in failed",
                    "This account does not have admin access. Ask an existing admin to grant it on the Users page, then try again.",
                ))
                .into_response();
            }
            // Set the dashboard session cookie and redirect to /admin/.
            // Same cookie shape /admin/login uses on success.
            let cookie = format!(
                "{name}={token}; HttpOnly; Path=/; SameSite=Strict; Max-Age={ttl}",
                name = SESSION_COOKIE,
                token = ok.token,
                ttl = state.cfg.session_ttl_secs,
            );
            let mut headers = HeaderMap::new();
            if let Ok(v) = HeaderValue::from_str(&cookie) {
                headers.insert(SET_COOKIE, v);
            }
            headers.insert(LOCATION, HeaderValue::from_static("/admin/"));
            (StatusCode::SEE_OTHER, headers).into_response()
        }
        Ok(_) => Html(html_page(
            "Sign-in complete",
            "You can close this window and return to RustDesk.",
        ))
        .into_response(),
        Err(msg) => Html(html_page("Sign-in failed", &html_escape(&msg))).into_response(),
    }
}

struct HandleOk {
    /// Bearer token freshly minted for the local user. For the admin flow
    /// we set it as `rd_admin_session`; for the desktop flow it's already
    /// stashed on the OidcSession row for `/api/oidc/auth-query` polling.
    token: String,
    user_is_admin: bool,
    is_admin_flow: bool,
}

async fn handle(state: Arc<AppState>, q: CallbackQuery) -> Result<HandleOk, String> {
    if q.state.is_empty() {
        return Err("missing state parameter".into());
    }
    let session = state
        .db
        .oidc_session_get_by_state(&q.state)
        .await
        .map_err(|e| e.to_string())?
        .ok_or_else(|| "unknown or expired oidc session (state)".to_string())?;

    if let Some(err) = q.error.as_deref().filter(|s| !s.is_empty()) {
        let detail = q
            .error_description
            .as_deref()
            .filter(|s| !s.is_empty())
            .unwrap_or(err);
        let _ = state
            .db
            .oidc_session_fail(&session.code, &format!("idp: {}", detail))
            .await;
        return Err(format!("identity provider returned an error: {}", detail));
    }
    if q.code.is_empty() {
        return Err("missing authorization code".into());
    }

    let provider = require_provider(&state, &session.provider)
        .await
        .map_err(|e| format!("{:?}", e))?;
    let disc = discovery::discover(&provider.issuer_url).await?;

    // Token exchange.
    let token_body = match discovery::http_post_form(
        &disc.token_endpoint,
        &[
            ("grant_type", "authorization_code"),
            ("code", &q.code),
            ("redirect_uri", &provider.redirect_url),
            ("client_id", &provider.client_id),
            ("client_secret", &provider.client_secret),
        ],
    )
    .await
    {
        Ok(b) => b,
        Err(e) => {
            let _ = state
                .db
                .oidc_session_fail(&session.code, &format!("token exchange: {}", e))
                .await;
            return Err(e);
        }
    };
    let token_resp: Value =
        serde_json::from_str(&token_body).map_err(|e| format!("parse token resp: {}", e))?;
    let access_token = token_resp
        .get("access_token")
        .and_then(|v| v.as_str())
        .ok_or_else(|| "token response missing access_token".to_string())?;

    // Fetch userinfo. We trust the userinfo endpoint as the authority on
    // the user's identity (sub + optional email + name).
    let userinfo_url = disc
        .userinfo_endpoint
        .as_deref()
        .ok_or_else(|| "provider has no userinfo_endpoint".to_string())?;
    let userinfo_body = discovery::http_get_with_bearer(userinfo_url, access_token).await?;
    let userinfo: Value = serde_json::from_str(&userinfo_body)
        .map_err(|e| format!("parse userinfo: {}", e))?;
    let sub = userinfo
        .get("sub")
        .and_then(|v| v.as_str())
        .ok_or_else(|| "userinfo missing sub".to_string())?;
    let email = userinfo.get("email").and_then(|v| v.as_str());
    let display_name = userinfo
        .get("name")
        .and_then(|v| v.as_str())
        .or_else(|| userinfo.get("preferred_username").and_then(|v| v.as_str()));

    // Optional role-based admin sync. When the provider is configured with
    // `admin_role`, we look up the userinfo claim at `roles_claim` (default
    // "roles") and set is_admin accordingly. Two shapes are supported:
    //   - object: presence of the role name as a key
    //     (Zitadel default: `"urn:zitadel:iam:org:project:roles":
    //       {"admin": {"<orgid>": "<orgname>"}}`)
    //   - array of strings: presence of the role name as an element
    //     (e.g. a custom claim mapping `"roles": ["admin", "user"]`)
    let desired_admin = provider.admin_role.as_deref().map(|role| {
        let claim_name = provider.roles_claim.as_deref().unwrap_or("roles");
        eval_admin_role(&userinfo, claim_name, role)
    });

    let user = state
        .db
        .user_upsert_oidc(sub, email, display_name, desired_admin)
        .await
        .map_err(|e| e.to_string())?;
    if user.status == 0 {
        return Err("user is disabled".into());
    }

    // Mint our own access token, store hashed, mark session complete.
    let token = mint_token();
    let sha = sha256_token(&token);
    let is_admin_flow = session.client_uuid == ADMIN_SENTINEL;
    // For admin-UI OIDC the "device id/uuid" fields carry the sentinel —
    // don't pollute the tokens.peer_* columns with it.
    let (token_peer_id, token_peer_uuid): (&str, &str) = if is_admin_flow {
        ("", "")
    } else {
        (
            session.client_id_str.as_str(),
            session.client_uuid.as_str(),
        )
    };
    state
        .db
        .token_insert(
            user.id,
            &sha,
            token_peer_id,
            token_peer_uuid,
            &session.device_info_json,
            state.cfg.session_ttl_secs,
        )
        .await
        .map_err(|e| e.to_string())?;
    // Best-effort device claim — same path as `/api/login`. Skipped for
    // admin-UI flow because the "device" is the operator's browser, not a
    // real RustDesk peer; calling device_claim with the sentinel would
    // insert a phantom row in device_sysinfo.
    if !is_admin_flow {
        state
            .db
            .device_claim(user.id, &session.client_id_str, &session.client_uuid)
            .await;
    }

    state
        .db
        .oidc_session_complete(&session.code, &token, user.id)
        .await
        .map_err(|e| e.to_string())?;
    Ok(HandleOk {
        token,
        user_is_admin: user.is_admin,
        is_admin_flow,
    })
}

fn html_page(title: &str, body: &str) -> String {
    format!(
        r#"<!doctype html>
<html><head><meta charset="utf-8"><title>{title}</title>
<style>
body {{ font-family: -apple-system, system-ui, sans-serif;
  display: flex; flex-direction: column; align-items: center;
  justify-content: center; height: 100vh; margin: 0;
  background: #0e0f12; color: #e6e6e6; }}
.card {{ background: #1c1e22; padding: 48px 56px;
  border-radius: 12px; box-shadow: 0 8px 24px rgba(0,0,0,.3); max-width: 480px; }}
h1 {{ margin: 0 0 16px; font-size: 22px; }}
p  {{ margin: 0; line-height: 1.5; color: #b8b8b8; }}
</style>
</head><body><div class="card">
<h1>{title}</h1>
<p>{body}</p>
</div></body></html>"#,
        title = title,
        body = body
    )
}

fn html_escape(s: &str) -> String {
    s.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
}

/// Returns true iff the userinfo's `claim_name` field carries `role` —
/// either as an object key (Zitadel) or as an element of a string array
/// (generic). Anything else (missing claim, wrong type, role not present)
/// is treated as "not admin" so a misconfigured claim demotes rather than
/// silently grants.
fn eval_admin_role(userinfo: &Value, claim_name: &str, role: &str) -> bool {
    let Some(node) = userinfo.get(claim_name) else {
        return false;
    };
    if let Some(obj) = node.as_object() {
        return obj.contains_key(role);
    }
    if let Some(arr) = node.as_array() {
        return arr
            .iter()
            .any(|v| v.as_str().map(|s| s == role).unwrap_or(false));
    }
    false
}

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

    #[test]
    fn admin_role_zitadel_object_shape() {
        let u = json!({
            "sub": "1",
            "urn:zitadel:iam:org:project:roles": {
                "admin": {"123": "myorg"},
                "user":  {"123": "myorg"},
            },
        });
        assert!(eval_admin_role(&u, "urn:zitadel:iam:org:project:roles", "admin"));
        assert!(!eval_admin_role(&u, "urn:zitadel:iam:org:project:roles", "owner"));
    }

    #[test]
    fn admin_role_generic_array_shape() {
        let u = json!({"sub": "1", "roles": ["admin", "user"]});
        assert!(eval_admin_role(&u, "roles", "admin"));
        assert!(!eval_admin_role(&u, "roles", "owner"));
    }

    #[test]
    fn admin_role_missing_claim_is_not_admin() {
        let u = json!({"sub": "1"});
        assert!(!eval_admin_role(&u, "roles", "admin"));
    }
}
