procyon 0.1.2

use std::path::{Path, PathBuf};

use rmcp::transport::auth::{AuthError, CredentialStore, StoredCredentials};

/// Loopback port for the authorization redirect. Fixed rather than ephemeral because the
/// redirect URI is registered with the authorization server during the flow, and a server may
/// pin it.
pub const REDIRECT_PORT: u16 = 8181;

pub fn redirect_uri() -> String {
    format!("http://127.0.0.1:{}/callback", REDIRECT_PORT)
}

fn credentials_dir() -> Result<PathBuf, String> {
    let base = dirs::data_dir().ok_or("Failed to locate a data directory")?;
    Ok(base.join("procyon").join("oauth"))
}

/// One file per server name, so several servers can be authorized independently.
pub fn credentials_path(server: &str) -> Result<PathBuf, String> {
    let safe: String = server
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
                c
            } else {
                '-'
            }
        })
        .collect();
    Ok(credentials_dir()?.join(format!("{}.json", safe)))
}

// The refresh token lives in here, so the file must not be readable by other local users.
#[cfg(unix)]
fn restrict_to_owner(path: &Path) -> std::io::Result<()> {
    use std::os::unix::fs::PermissionsExt;
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
}

#[cfg(not(unix))]
fn restrict_to_owner(_path: &Path) -> std::io::Result<()> {
    Ok(())
}

/// Persists OAuth credentials so a browser sign-in survives a restart. Failures are surfaced as
/// `AuthError` rather than swallowed: silently losing a refresh token would push the user through
/// the browser again with no explanation.
pub struct FileCredentialStore {
    path: PathBuf,
}

impl FileCredentialStore {
    pub fn new(path: PathBuf) -> Self {
        Self { path }
    }
}

#[async_trait::async_trait]
impl CredentialStore for FileCredentialStore {
    async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
        let raw = match tokio::fs::read_to_string(&self.path).await {
            Ok(raw) => raw,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
            Err(e) => {
                return Err(AuthError::InternalError(format!(
                    "Failed to read {}: {}",
                    self.path.display(),
                    e
                )))
            }
        };

        // A corrupt file must not be fatal: dropping it costs one browser sign-in, whereas
        // failing here would leave the server permanently unusable.
        Ok(serde_json::from_str(&raw).ok())
    }

    async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> {
        if let Some(parent) = self.path.parent() {
            tokio::fs::create_dir_all(parent).await.map_err(|e| {
                AuthError::InternalError(format!("Failed to create {}: {}", parent.display(), e))
            })?;
        }

        let json = serde_json::to_string_pretty(&credentials).map_err(|e| {
            AuthError::InternalError(format!("Failed to encode credentials: {}", e))
        })?;

        tokio::fs::write(&self.path, json).await.map_err(|e| {
            AuthError::InternalError(format!("Failed to write {}: {}", self.path.display(), e))
        })?;

        restrict_to_owner(&self.path).map_err(|e| {
            AuthError::InternalError(format!(
                "Failed to restrict permissions on {}: {}",
                self.path.display(),
                e
            ))
        })
    }

    async fn clear(&self) -> Result<(), AuthError> {
        match tokio::fs::remove_file(&self.path).await {
            Ok(()) => Ok(()),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
            Err(e) => Err(AuthError::InternalError(format!(
                "Failed to remove {}: {}",
                self.path.display(),
                e
            ))),
        }
    }
}

/// When the stored access token runs out, for warning the user before a silent refresh.
pub fn expires_at(credentials: &StoredCredentials) -> Option<u64> {
    use oauth2::TokenResponse;
    let received = credentials.token_received_at?;
    let lifetime = credentials.token_response.as_ref()?.expires_in()?;
    Some(received + lifetime.as_secs())
}

#[derive(Debug, PartialEq)]
pub struct Callback {
    pub code: String,
    pub state: String,
    /// The absolute redirect URL, so the SDK can also read `iss` and perform RFC 9207 issuer
    /// validation. Parsing only `code` and `state` here made the token exchange fail against a
    /// server that requires it.
    pub url: String,
}

/// Parses the query of the redirect the authorization server sends to the loopback listener.
/// An `error` parameter is reported rather than being mistaken for a missing code.
pub fn parse_callback_query(query: &str) -> Result<Callback, String> {
    let mut code = None;
    let mut state = None;
    let mut error = None;
    let mut description = None;

    for pair in query.split('&') {
        let Some((key, value)) = pair.split_once('=') else {
            continue;
        };
        let value = percent_decode(value);
        match key {
            "code" => code = Some(value),
            "state" => state = Some(value),
            "error" => error = Some(value),
            "error_description" => description = Some(value),
            _ => {}
        }
    }

    if let Some(error) = error {
        return Err(match description {
            Some(description) => format!("{}: {}", error, description),
            None => error,
        });
    }

    match (code, state) {
        (Some(code), Some(state)) => Ok(Callback {
            code,
            state,
            url: format!("{}?{}", redirect_uri(), query),
        }),
        (None, _) => Err("The redirect carried no authorization code".to_string()),
        (_, None) => Err("The redirect carried no state parameter".to_string()),
    }
}

fn percent_decode(value: &str) -> String {
    let bytes = value.as_bytes();
    let mut out = Vec::with_capacity(bytes.len());
    let mut i = 0;

    while i < bytes.len() {
        match bytes[i] {
            b'+' => {
                out.push(b' ');
                i += 1;
            }
            b'%' if i + 2 < bytes.len() => {
                match u8::from_str_radix(&value[i + 1..i + 3], 16) {
                    Ok(byte) => {
                        out.push(byte);
                        i += 3;
                    }
                    // Not a valid escape; keep it literal rather than losing the character.
                    Err(_) => {
                        out.push(bytes[i]);
                        i += 1;
                    }
                }
            }
            byte => {
                out.push(byte);
                i += 1;
            }
        }
    }

    String::from_utf8_lossy(&out).to_string()
}

const BROWSER_RESPONSE_OK: &str = "HTTP/1.1 200 OK\r\n\
     Content-Type: text/html; charset=utf-8\r\n\
     Connection: close\r\n\r\n\
     <html><body style=\"font-family:system-ui;padding:3rem\">\
     <h2>Procyon is authorized</h2><p>You can close this tab and return to the terminal.</p>\
     </body></html>";

const BROWSER_RESPONSE_ERR: &str = "HTTP/1.1 400 Bad Request\r\n\
     Content-Type: text/html; charset=utf-8\r\n\
     Connection: close\r\n\r\n\
     <html><body style=\"font-family:system-ui;padding:3rem\">\
     <h2>Authorization failed</h2><p>Return to the terminal for the reason.</p>\
     </body></html>";

/// Reads the request line of the redirect and returns its query string. Bound to loopback only,
/// so nothing off the machine can reach it.
pub async fn wait_for_redirect(
    listener: tokio::net::TcpListener,
    timeout: std::time::Duration,
) -> Result<Callback, String> {
    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};

    let accept = tokio::time::timeout(timeout, listener.accept()).await;
    let (mut stream, _) = match accept {
        Ok(Ok(pair)) => pair,
        Ok(Err(e)) => return Err(format!("Failed to accept the redirect: {}", e)),
        Err(_) => {
            return Err(format!(
                "No redirect arrived within {}s. Authorization was not completed.",
                timeout.as_secs()
            ))
        }
    };

    let mut request_line = String::new();
    BufReader::new(&mut stream)
        .read_line(&mut request_line)
        .await
        .map_err(|e| format!("Failed to read the redirect: {}", e))?;

    // "GET /callback?code=...&state=... HTTP/1.1"
    let target = request_line.split_whitespace().nth(1).unwrap_or("");
    let query = target.split_once('?').map(|(_, q)| q).unwrap_or("");
    let parsed = parse_callback_query(query);

    let body = if parsed.is_ok() {
        BROWSER_RESPONSE_OK
    } else {
        BROWSER_RESPONSE_ERR
    };
    let _ = stream.write_all(body.as_bytes()).await;
    let _ = stream.shutdown().await;

    parsed
}

/// Binds the loopback listener up front, so a port already in use is reported before the user is
/// sent to a browser that would redirect nowhere.
pub async fn bind_redirect_listener() -> Result<tokio::net::TcpListener, String> {
    tokio::net::TcpListener::bind(("127.0.0.1", REDIRECT_PORT))
        .await
        .map_err(|e| {
            format!(
                "Cannot listen on 127.0.0.1:{} for the OAuth redirect: {}",
                REDIRECT_PORT, e
            )
        })
}

/// Best effort: if no browser can be launched the caller still shows the URL, so the flow is
/// completable by hand.
pub async fn open_browser(url: &str) -> bool {
    let opener = if cfg!(target_os = "macos") {
        "open"
    } else if cfg!(target_os = "windows") {
        "explorer"
    } else {
        "xdg-open"
    };

    tokio::process::Command::new(opener)
        .arg(url)
        // The TUI owns the terminal; a launcher writing to it would corrupt the display.
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status()
        .await
        .map(|status| status.success())
        .unwrap_or(false)
}

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

    #[test]
    fn parses_a_normal_redirect() {
        let callback = parse_callback_query("code=abc123&state=xyz789").unwrap();
        assert_eq!(callback.code, "abc123");
        assert_eq!(callback.state, "xyz789");
    }

    #[test]
    fn ignores_extra_parameters() {
        let callback =
            parse_callback_query("iss=https%3A%2F%2Fx&code=a&state=b&scope=read").unwrap();
        assert_eq!(callback.code, "a");
        assert_eq!(callback.state, "b");
    }

    #[test]
    fn percent_and_plus_escapes_are_decoded() {
        let callback = parse_callback_query("code=a%2Fb%2Bc&state=x+y").unwrap();
        assert_eq!(callback.code, "a/b+c");
        assert_eq!(callback.state, "x y");
    }

    #[test]
    fn an_error_redirect_reports_the_server_reason() {
        let err = parse_callback_query("error=access_denied&error_description=User%20said%20no")
            .unwrap_err();
        assert!(err.contains("access_denied"), "got {}", err);
        assert!(err.contains("User said no"), "got {}", err);
    }

    #[test]
    fn an_error_without_a_description_still_reports() {
        let err = parse_callback_query("error=server_error").unwrap_err();
        assert_eq!(err, "server_error");
    }

    #[test]
    fn a_missing_code_is_distinguished_from_a_missing_state() {
        assert!(parse_callback_query("state=x")
            .unwrap_err()
            .contains("no authorization code"));
        assert!(parse_callback_query("code=x")
            .unwrap_err()
            .contains("no state parameter"));
    }

    #[test]
    fn an_empty_query_does_not_panic() {
        assert!(parse_callback_query("").is_err());
    }

    #[test]
    fn a_malformed_escape_is_kept_literal() {
        // Truncated escapes must not silently delete characters from a code.
        assert_eq!(percent_decode("a%zzb"), "a%zzb");
        assert_eq!(percent_decode("trailing%"), "trailing%");
    }

    #[test]
    fn the_redirect_uri_is_loopback_only() {
        let uri = redirect_uri();
        assert!(uri.starts_with("http://127.0.0.1:"), "got {}", uri);
        assert!(uri.ends_with("/callback"));
    }

    #[test]
    fn credentials_are_kept_per_server() {
        let a = credentials_path("raven").unwrap();
        let b = credentials_path("other").unwrap();
        assert_ne!(a, b);
        assert!(a.to_string_lossy().ends_with("raven.json"));
    }

    #[test]
    fn a_server_name_cannot_escape_the_credentials_directory() {
        let path = credentials_path("../../etc/shadow").unwrap();
        let name = path.file_name().unwrap().to_string_lossy().to_string();
        assert!(!name.contains('/'), "got {}", name);
        assert_eq!(path.parent(), Some(credentials_dir().unwrap().as_path()));
    }

    #[test]
    fn expiry_is_unknown_before_a_token_is_issued() {
        // A dynamically registered client exists before any token does.
        let credentials = StoredCredentials::new("client-abc".to_string(), None, Vec::new(), None);
        assert!(expires_at(&credentials).is_none());
    }

    #[tokio::test]
    async fn a_missing_credentials_file_loads_as_none() {
        let temp = tempfile::tempdir().unwrap();
        let store = FileCredentialStore::new(temp.path().join("absent.json"));
        assert!(store.load().await.unwrap().is_none());
    }

    #[tokio::test]
    async fn a_corrupt_credentials_file_loads_as_none_rather_than_failing() {
        let temp = tempfile::tempdir().unwrap();
        let path = temp.path().join("broken.json");
        tokio::fs::write(&path, "{ not json").await.unwrap();

        let store = FileCredentialStore::new(path);
        assert!(
            store.load().await.unwrap().is_none(),
            "a corrupt file must cost a sign-in, not break the server"
        );
    }

    #[tokio::test]
    async fn clearing_an_absent_file_succeeds() {
        let temp = tempfile::tempdir().unwrap();
        let store = FileCredentialStore::new(temp.path().join("absent.json"));
        assert!(store.clear().await.is_ok());
    }

    #[tokio::test]
    async fn saved_credentials_round_trip_and_are_owner_only() {
        let temp = tempfile::tempdir().unwrap();
        let path = temp.path().join("creds.json");
        let store = FileCredentialStore::new(path.clone());

        let credentials = StoredCredentials::new("client-abc".to_string(), None, Vec::new(), None);
        store.save(credentials).await.unwrap();

        let loaded = store.load().await.unwrap().expect("credentials");
        assert_eq!(loaded.client_id, "client-abc");

        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mode = tokio::fs::metadata(&path)
                .await
                .unwrap()
                .permissions()
                .mode();
            assert_eq!(
                mode & 0o777,
                0o600,
                "the refresh token must not be world readable"
            );
        }
    }

    #[tokio::test]
    async fn a_redirect_that_never_arrives_times_out_with_a_clear_message() {
        let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
            .await
            .unwrap();
        let err = wait_for_redirect(listener, std::time::Duration::from_millis(50))
            .await
            .unwrap_err();
        assert!(err.contains("No redirect arrived"), "got {}", err);
    }

    #[tokio::test]
    async fn a_real_redirect_is_read_from_the_socket() {
        use tokio::io::AsyncWriteExt;

        let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
            .await
            .unwrap();
        let port = listener.local_addr().unwrap().port();

        let server = tokio::spawn(async move {
            wait_for_redirect(listener, std::time::Duration::from_secs(5)).await
        });

        let mut client = tokio::net::TcpStream::connect(("127.0.0.1", port))
            .await
            .unwrap();
        client
            .write_all(b"GET /callback?code=THECODE&state=THESTATE HTTP/1.1\r\nHost: x\r\n\r\n")
            .await
            .unwrap();

        let callback = server.await.unwrap().unwrap();
        assert_eq!(callback.code, "THECODE");
        assert_eq!(callback.state, "THESTATE");
    }

    #[tokio::test]
    async fn a_redirect_carrying_an_error_is_surfaced() {
        use tokio::io::AsyncWriteExt;

        let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
            .await
            .unwrap();
        let port = listener.local_addr().unwrap().port();

        let server = tokio::spawn(async move {
            wait_for_redirect(listener, std::time::Duration::from_secs(5)).await
        });

        let mut client = tokio::net::TcpStream::connect(("127.0.0.1", port))
            .await
            .unwrap();
        client
            .write_all(b"GET /callback?error=access_denied HTTP/1.1\r\n\r\n")
            .await
            .unwrap();

        let err = server.await.unwrap().unwrap_err();
        assert!(err.contains("access_denied"), "got {}", err);
    }
}