//! `mcptest login` interactive OAuth 2.1 flow.
//!
//! Spawns a loopback HTTP server on an ephemeral port, opens the
//! operator's browser to the authorization URL, receives the redirect,
//! exchanges the code for tokens, and writes them into the existing
//! [`mcptest_core::auth::TokenCache`]. The `--no-browser` flag skips the
//! browser open and prints the URL so headless environments (CI, SSH
//! sessions, integration tests) can paste it into a real browser. The
//! `--device` flag (WOR-1286) swaps the loopback dance for the RFC 8628
//! device-code grant; that path lives in the sibling `login_device`
//! module and shares discovery, DCR, and the cache-write tail with this
//! one.

use std::io::{BufRead, BufReader, Write};
use std::net::TcpListener;
use std::path::PathBuf;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use anyhow::{anyhow, Context, Result};
use mcptest_core::auth::{
    exchange_code, register_client, AuthorizationUrl, ClientCache, PkceParams,
    RegisterClientRequest, Token, TokenCache, TokenExchangeRequest,
};
use serde::Deserialize;

/// Plain-data inputs to [`run_login`]. The CLI layer builds one of these
/// from its `LoginArgs`.
#[derive(Debug, Clone)]
pub struct LoginRunArgs {
    /// Server URL the operator is authenticating against. Required; the
    /// caller is responsible for falling back to the YAML config when
    /// `--url` is omitted.
    pub server_url: String,
    /// OAuth `client_id` to use when DCR is not available. When DCR is
    /// available this acts as a fallback that the runner falls through
    /// to only if registration fails.
    pub client_id: Option<String>,
    /// Optional override for the YAML config the loader otherwise picks
    /// up from `--config` or `mcptest.yml`. Used by tests that point at
    /// a fixture path.
    pub config_path: Option<PathBuf>,
    /// When true, skip opening the browser and print the auth URL
    /// instead. The flow still spins up the loopback listener so the
    /// operator can finish the dance manually.
    pub no_browser: bool,
    /// When true, run the RFC 8628 device-code grant (WOR-1286) instead
    /// of the browser dance: print a verification URL plus user code and
    /// poll the token endpoint until a human approves on another device.
    pub device: bool,
    /// When true, evict any cached token and cached DCR client metadata
    /// for `server_url` before running the flow. The operator uses
    /// this when an IdP rotates a client or revokes a session.
    pub force: bool,
    /// Override for the token cache root. `None` means default to
    /// `~/.mcptest/auth/`.
    pub cache_root: Option<PathBuf>,
    /// Timeout for the user to complete the browser dance. Defaults to
    /// five minutes.
    pub timeout: Duration,
}

impl Default for LoginRunArgs {
    fn default() -> Self {
        Self {
            server_url: String::new(),
            client_id: None,
            config_path: None,
            no_browser: false,
            device: false,
            force: false,
            cache_root: None,
            timeout: Duration::from_secs(300),
        }
    }
}

/// OAuth 2.1 server metadata (RFC 8414 subset). Only the fields the
/// login flow consumes are parsed; everything else is ignored so a
/// stricter `deny_unknown_fields` does not bite us on benign extras.
/// `pub(super)` so the device-code sibling module shares discovery.
#[derive(Debug, Deserialize)]
pub(super) struct ServerMetadata {
    /// Authorization endpoint the browser flow redirects to.
    pub(super) authorization_endpoint: String,
    /// Token endpoint both grants exchange against.
    pub(super) token_endpoint: String,
    /// Optional RFC 7591 dynamic-registration endpoint.
    #[serde(default)]
    pub(super) registration_endpoint: Option<String>,
    /// Optional RFC 8628 device authorization endpoint (WOR-1286).
    /// Absent means the IdP does not support the device grant.
    #[serde(default)]
    pub(super) device_authorization_endpoint: Option<String>,
    /// Scopes the server advertises; joined into the scope request.
    #[serde(default)]
    pub(super) scopes_supported: Option<Vec<String>>,
}

/// Run the interactive login flow. Returns the process exit code.
pub fn run_login(args: LoginRunArgs) -> Result<i32> {
    if args.server_url.is_empty() {
        return Err(anyhow!(
            "mcptest login requires --url or a server URL in mcptest.yml"
        ));
    }
    tracing::info!(
        target: "mcptest::commands::login",
        server_url = %args.server_url,
        force = args.force,
        "login start"
    );

    // Clear cached artifacts up front when --force is set so a stale
    // entry cannot survive even when the subsequent metadata fetch or
    // token exchange fails.
    if args.force {
        force_clear_cache(&args)?;
    }

    let metadata = fetch_metadata(&args.server_url).context("fetch OAuth server metadata")?;

    // WOR-1286: the device grant shares discovery, caching, and DCR
    // with the browser flow but needs no loopback listener, so branch
    // before binding one.
    if args.device {
        return super::login_device::run_device_login(&args, &metadata);
    }

    // Loopback listener on an ephemeral port. We pick the port now so
    // the redirect URI is stable across DCR, the authorization URL,
    // and the eventual code exchange.
    let listener =
        TcpListener::bind("127.0.0.1:0").context("could not bind 127.0.0.1 loopback listener")?;
    let port = listener
        .local_addr()
        .context("local addr from loopback listener")?
        .port();
    let redirect_uri = format!("http://127.0.0.1:{port}/callback");

    // Cache layout: tokens and DCR results sit under the same root,
    // keyed by sha256(server_url). Tests override `cache_root`.
    let token_cache = resolve_token_cache(&args)?;

    let registration = RegisterClientRequest::public_client("mcptest", &redirect_uri);
    let client_id = resolve_client_id(
        &token_cache,
        &metadata,
        &args.server_url,
        &registration,
        args.client_id.clone(),
    )?;

    let pkce = PkceParams::generate();
    let state = mcptest_core::auth::random_state();
    build_and_open_authorization_url(
        &metadata,
        &client_id,
        &redirect_uri,
        scope_of(&metadata),
        &pkce,
        &state,
        args.no_browser,
    )?;

    let now_unix = now_unix();
    let token = await_callback_and_exchange(
        &listener,
        &metadata,
        &client_id,
        &redirect_uri,
        pkce,
        &state,
        args.timeout,
        now_unix,
    )?;
    finish_login(
        &args,
        &token_cache,
        &metadata.token_endpoint,
        client_id,
        token,
        now_unix,
    )
}

/// Block on the loopback callback, enforce the CSRF `state` echo, and
/// exchange the authorization code for tokens. Split out of
/// [`run_login`] so the top-level flow stays readable at a glance.
#[allow(clippy::too_many_arguments)]
fn await_callback_and_exchange(
    listener: &TcpListener,
    metadata: &ServerMetadata,
    client_id: &str,
    redirect_uri: &str,
    pkce: PkceParams,
    state: &str,
    timeout: Duration,
    now_unix: u64,
) -> Result<Token> {
    listener
        .set_nonblocking(false)
        .context("set listener to blocking mode")?;
    let callback = wait_for_callback(listener, timeout)?;
    if callback.state != state {
        return Err(anyhow!(
            "state mismatch on OAuth callback: expected `{state}`, got `{}` (possible CSRF)",
            callback.state
        ));
    }
    exchange_code(
        &TokenExchangeRequest {
            token_endpoint: metadata.token_endpoint.clone(),
            client_id: client_id.to_owned(),
            client_secret: None,
            code: callback.code,
            redirect_uri: redirect_uri.to_owned(),
            code_verifier: pkce.code_verifier,
        },
        now_unix,
    )
    .context("exchange authorization code for tokens")
}

/// Stamp refresh metadata onto the fresh token, persist it, and print
/// the success line. Shared tail of the browser (PKCE) and device-code
/// (WOR-1286) paths so both produce identical cache entries and
/// operator output.
pub(super) fn finish_login(
    args: &LoginRunArgs,
    token_cache: &TokenCache,
    token_endpoint: &str,
    client_id: String,
    mut token: Token,
    now_unix: u64,
) -> Result<i32> {
    // Stamp the endpoint + client id so a later `mcptest run` can refresh this
    // token in place when it expires, instead of forcing a re-login.
    token.token_endpoint = Some(token_endpoint.to_owned());
    token.client_id = Some(client_id);

    token_cache
        .write(&args.server_url, &token)
        .context("write token cache")?;
    let lifetime = token.expires_at_unix.saturating_sub(now_unix);
    tracing::info!(
        target: "mcptest::commands::login",
        server_url = %args.server_url,
        lifetime_seconds = lifetime,
        "login success"
    );
    println!(
        "login successful; token cached for {} (expires in {} seconds)",
        args.server_url, lifetime
    );
    Ok(0)
}

/// Space-join `scopes_supported` into the scope request string. Empty
/// means "request no scope"; the downstream builders omit the parameter.
pub(super) fn scope_of(metadata: &ServerMetadata) -> String {
    metadata
        .scopes_supported
        .as_ref()
        .map(|s| s.join(" "))
        .unwrap_or_default()
}

/// Wall-clock seconds since the UNIX epoch. Clamped to 0 on a clock set
/// before 1970 (the only way `duration_since` fails) so token expiry
/// math never panics in a release binary.
pub(super) fn now_unix() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

/// Resolve the token cache from `--cache-root` or the user default.
pub(super) fn resolve_token_cache(args: &LoginRunArgs) -> Result<TokenCache> {
    match &args.cache_root {
        Some(root) => Ok(TokenCache::with_root(root.clone())),
        None => TokenCache::user_default().context("resolve default token cache root"),
    }
}

/// Clear the cached token and client metadata for the server. Used by the
/// `--force` path so a stale entry cannot survive a later failure.
fn force_clear_cache(args: &LoginRunArgs) -> Result<()> {
    let cache = resolve_token_cache(args)?;
    clear_cached_artifacts(&cache, &args.server_url)
        .context("--force: clear cached token and client metadata")
}

/// Build the authorization URL and either open it in the browser or print it
/// for the operator to open manually (`--no-browser`).
#[allow(clippy::too_many_arguments)]
fn build_and_open_authorization_url(
    metadata: &ServerMetadata,
    client_id: &str,
    redirect_uri: &str,
    scope: String,
    pkce: &PkceParams,
    state: &str,
    no_browser: bool,
) -> Result<String> {
    let url = AuthorizationUrl {
        authorization_endpoint: metadata.authorization_endpoint.clone(),
        client_id: client_id.to_owned(),
        redirect_uri: redirect_uri.to_owned(),
        scope,
        state: state.to_owned(),
        code_challenge: pkce.code_challenge.clone(),
        code_challenge_method: pkce.code_challenge_method,
    }
    .build()
    .context("build authorization URL")?;

    if no_browser {
        println!("Open this URL in your browser to continue:\n  {url}");
    } else {
        println!("Opening browser to:\n  {url}");
        if let Err(err) = webbrowser::open(&url) {
            eprintln!("warning: could not open browser ({err}); copy this URL manually:\n  {url}");
        }
    }
    Ok(url)
}

/// Resolve a `client_id` for the login flow. Cached DCR results win;
/// otherwise, when the server advertises a `registration_endpoint`, run
/// DCR with the grant-specific `registration` request (caching the
/// result for later runs); otherwise fall back to the operator-provided
/// `--client-id`. Shared by the PKCE and device-code (WOR-1286) paths,
/// which differ only in the registration request they build.
pub(super) fn resolve_client_id(
    token_cache: &TokenCache,
    metadata: &ServerMetadata,
    server_url: &str,
    registration: &RegisterClientRequest,
    fallback: Option<String>,
) -> Result<String> {
    let client_cache = ClientCache::new(token_cache);
    if let Some(cached) = client_cache.read(server_url).context("read client cache")? {
        return Ok(cached.client_id);
    }
    if let Some(endpoint) = metadata.registration_endpoint.as_deref() {
        let response =
            register_client(endpoint, registration).context("dynamic client registration")?;
        client_cache
            .write(server_url, &response)
            .context("persist DCR client metadata")?;
        return Ok(response.client_id);
    }
    fallback.ok_or_else(|| {
        anyhow!(
            "server does not advertise a registration_endpoint; pass --client-id <ID> on the CLI"
        )
    })
}

/// Remove the cached token JSON and the DCR client metadata for
/// `server_url`. The lock file is left alone; a stale lock is harmless
/// across runs. Missing files are not an error. Used by `--force`.
fn clear_cached_artifacts(token_cache: &TokenCache, server_url: &str) -> Result<()> {
    let key = token_cache.key_for(server_url);
    let root = token_cache.root().to_path_buf();
    for suffix in ["json", "client.json"] {
        let path = root.join(format!("{key}.{suffix}"));
        match std::fs::remove_file(&path) {
            Ok(()) => {}
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
            Err(e) => return Err(anyhow!("remove {}: {e}", path.display())),
        }
    }
    Ok(())
}

/// Fetch and parse `.well-known/oauth-authorization-server` from the
/// given server URL. The URL may be the bare origin
/// (`https://idp.example`) or include the well-known path already
/// (tests point at a wiremock root and pass the suffix in `--url`).
/// RFC 9728 protected-resource metadata. We only read the field that names
/// the authorization server(s) guarding the resource.
#[derive(Debug, Deserialize)]
struct ProtectedResource {
    #[serde(default)]
    authorization_servers: Vec<String>,
}

/// Origin (`scheme://host[:port]`) of a URL. Well-known metadata documents
/// live at the origin root per RFC 8615, not under the resource path, so we
/// strip the path before appending `.well-known/...`.
fn origin_of(raw: &str) -> Result<String> {
    let u = url::Url::parse(raw).with_context(|| format!("parse URL {raw}"))?;
    let scheme = u.scheme();
    let host = u.host_str().ok_or_else(|| anyhow!("{raw} has no host"))?;
    Ok(match u.port() {
        Some(p) => format!("{scheme}://{host}:{p}"),
        None => format!("{scheme}://{host}"),
    })
}

/// Discover the OAuth metadata for an MCP server (RFC 9728 + RFC 8414).
///
/// 1. If the caller passed a metadata URL directly, fetch it as-is.
/// 2. Otherwise fetch the protected-resource metadata at the server's origin
///    (`/.well-known/oauth-protected-resource`) and follow its
///    `authorization_servers` to the authorization server's metadata. This is
///    how an MCP client reaches an auth server hosted on a different origin.
/// 3. If the resource publishes no protected-resource document, fall back to
///    looking for authorization-server metadata at the server's own origin.
fn fetch_metadata(server_url: &str) -> Result<ServerMetadata> {
    let client = reqwest::blocking::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()
        .context("reqwest client")?;

    if server_url.contains("/.well-known/oauth-authorization-server") {
        return fetch_as_metadata(&client, server_url);
    }

    let server_origin = origin_of(server_url)?;
    let mut as_origins: Vec<String> = Vec::new();
    if let Some(origin) = discover_auth_server(&client, &server_origin) {
        as_origins.push(origin);
    }
    // Always keep the server's own origin as a fallback: many servers host
    // their authorization endpoints on the same origin and skip the
    // protected-resource document.
    if !as_origins.contains(&server_origin) {
        as_origins.push(server_origin);
    }

    let mut last_err = None;
    for origin in as_origins {
        let meta_url = format!(
            "{}/.well-known/oauth-authorization-server",
            origin.trim_end_matches('/')
        );
        match fetch_as_metadata(&client, &meta_url) {
            Ok(meta) => return Ok(meta),
            Err(e) => last_err = Some(e),
        }
    }
    Err(last_err.unwrap_or_else(|| anyhow!("could not discover OAuth metadata for {server_url}")))
}

/// Fetch the RFC 9728 protected-resource document at `server_origin` and return
/// the origin of its first authorization server, if any. A missing document is
/// not an error: the caller falls back to the server's own origin.
fn discover_auth_server(client: &reqwest::blocking::Client, server_origin: &str) -> Option<String> {
    let url = format!(
        "{}/.well-known/oauth-protected-resource",
        server_origin.trim_end_matches('/')
    );
    let resp = client.get(&url).send().ok()?;
    if !resp.status().is_success() {
        return None;
    }
    let doc: ProtectedResource = resp.json().ok()?;
    let first = doc.authorization_servers.into_iter().next()?;
    origin_of(&first).ok()
}

/// GET and parse authorization-server metadata (RFC 8414) from a metadata URL.
fn fetch_as_metadata(client: &reqwest::blocking::Client, url: &str) -> Result<ServerMetadata> {
    let resp = client.get(url).send().context("GET metadata")?;
    let status = resp.status();
    if !status.is_success() {
        return Err(anyhow!("metadata fetch failed: HTTP {status} from {url}"));
    }
    let body = resp.text().context("read metadata body")?;
    serde_json::from_str::<ServerMetadata>(&body)
        .with_context(|| format!("parse OAuth metadata from {url}"))
}

/// Result of a successful loopback callback.
#[derive(Debug)]
struct Callback {
    code: String,
    state: String,
}

/// Block on the loopback listener until the IdP completes the redirect
/// or `timeout` elapses. Returns the parsed `code` and `state` query
/// parameters from the first request whose path starts with
/// `/callback`.
fn wait_for_callback(listener: &TcpListener, timeout: Duration) -> Result<Callback> {
    let deadline = Instant::now() + timeout;
    listener
        .set_nonblocking(true)
        .context("set listener nonblocking")?;
    loop {
        if Instant::now() >= deadline {
            return Err(anyhow!(
                "timed out waiting for OAuth callback after {:?}",
                timeout
            ));
        }
        match listener.accept() {
            Ok((mut stream, _)) => {
                stream
                    .set_read_timeout(Some(Duration::from_secs(2)))
                    .context("read timeout on callback stream")?;
                stream
                    .set_write_timeout(Some(Duration::from_secs(2)))
                    .context("write timeout on callback stream")?;
                let request_line = read_request_line(&mut stream)?;
                let response = b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nConnection: close\r\nContent-Length: 32\r\n\r\nmcptest login complete. Bye!\n";
                let _ = stream.write_all(response);
                let _ = stream.flush();
                return parse_callback(&request_line);
            }
            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                std::thread::sleep(Duration::from_millis(50));
            }
            Err(e) => return Err(anyhow!("loopback accept failed: {e}")),
        }
    }
}

/// Pull the first line ("GET /callback?... HTTP/1.1") off the wire.
/// We do not parse the rest of the request; the redirect carries
/// everything we need in the URL.
fn read_request_line(stream: &mut std::net::TcpStream) -> Result<String> {
    let mut reader = BufReader::new(stream);
    let mut line = String::new();
    reader
        .read_line(&mut line)
        .context("read callback request line")?;
    Ok(line)
}

fn parse_callback(request_line: &str) -> Result<Callback> {
    // "GET /callback?code=...&state=... HTTP/1.1\r\n"
    let path_part = request_line
        .split_whitespace()
        .nth(1)
        .ok_or_else(|| anyhow!("malformed HTTP request line: {request_line:?}"))?;
    let parsed = url::Url::parse(&format!("http://127.0.0.1{path_part}"))
        .with_context(|| format!("parse callback path: {path_part}"))?;

    let mut code: Option<String> = None;
    let mut state: Option<String> = None;
    let mut error: Option<String> = None;
    let mut error_description: Option<String> = None;
    for (k, v) in parsed.query_pairs() {
        match k.as_ref() {
            "code" => code = Some(v.into_owned()),
            "state" => state = Some(v.into_owned()),
            "error" => error = Some(v.into_owned()),
            "error_description" => error_description = Some(v.into_owned()),
            _ => {}
        }
    }
    if let Some(err) = error {
        let desc = error_description.unwrap_or_default();
        return Err(anyhow!("IdP returned OAuth error `{err}`: {desc}"));
    }
    Ok(Callback {
        code: code.ok_or_else(|| anyhow!("callback missing `code` parameter"))?,
        state: state.ok_or_else(|| anyhow!("callback missing `state` parameter"))?,
    })
}

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

    #[test]
    fn parse_callback_extracts_code_and_state() {
        let cb = parse_callback("GET /callback?code=abc&state=xyz HTTP/1.1\r\n").unwrap();
        assert_eq!(cb.code, "abc");
        assert_eq!(cb.state, "xyz");
    }

    #[test]
    fn parse_callback_surfaces_oauth_error() {
        let err = parse_callback(
            "GET /callback?error=access_denied&error_description=user+said+no HTTP/1.1\r\n",
        )
        .unwrap_err();
        assert!(err.to_string().contains("access_denied"));
    }

    #[test]
    fn parse_callback_rejects_missing_code() {
        let err = parse_callback("GET /callback?state=xyz HTTP/1.1\r\n").unwrap_err();
        assert!(err.to_string().contains("missing `code`"));
    }

    #[test]
    fn origin_of_strips_the_resource_path() {
        // Well-known docs live at the origin root (RFC 8615), not under the
        // resource path, so the path must be dropped.
        assert_eq!(
            origin_of("https://mcp.example.com/secure/mcp").unwrap(),
            "https://mcp.example.com"
        );
        assert_eq!(
            origin_of("http://localhost:8799/secure/mcp").unwrap(),
            "http://localhost:8799"
        );
        assert_eq!(
            origin_of("https://auth.example.com").unwrap(),
            "https://auth.example.com"
        );
    }

    #[tokio::test]
    async fn fetch_metadata_follows_protected_resource_to_the_auth_server() {
        use serde_json::json;
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        let server = MockServer::start().await;
        let base = server.uri();
        // RFC 9728: the resource points at its authorization server.
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-protected-resource"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "resource": format!("{base}/secure/mcp"),
                "authorization_servers": [base],
            })))
            .mount(&server)
            .await;
        // RFC 8414: the authorization server's metadata, at the origin root.
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "authorization_endpoint": format!("{base}/authorize"),
                "token_endpoint": format!("{base}/token"),
                "registration_endpoint": format!("{base}/register"),
            })))
            .mount(&server)
            .await;

        let server_url = format!("{base}/secure/mcp");
        let meta = tokio::task::spawn_blocking(move || fetch_metadata(&server_url))
            .await
            .unwrap()
            .expect("metadata discovered via protected-resource");
        assert_eq!(meta.authorization_endpoint, format!("{base}/authorize"));
        assert_eq!(meta.token_endpoint, format!("{base}/token"));
        assert!(meta.registration_endpoint.is_some());
    }

    #[tokio::test]
    async fn fetch_metadata_falls_back_to_the_server_origin() {
        use serde_json::json;
        use wiremock::matchers::{method, path};
        use wiremock::{Mock, MockServer, ResponseTemplate};

        // No protected-resource document: the 404 makes discovery fall back to
        // looking for authorization-server metadata at the server's own origin.
        let server = MockServer::start().await;
        let base = server.uri();
        Mock::given(method("GET"))
            .and(path("/.well-known/oauth-authorization-server"))
            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
                "authorization_endpoint": format!("{base}/authorize"),
                "token_endpoint": format!("{base}/token"),
            })))
            .mount(&server)
            .await;

        let server_url = format!("{base}/secure/mcp");
        let meta = tokio::task::spawn_blocking(move || fetch_metadata(&server_url))
            .await
            .unwrap()
            .expect("metadata found at the server origin");
        assert_eq!(meta.token_endpoint, format!("{base}/token"));
    }
}
