linkshieldai 0.2.0

use std::io::{Read, Write};
use std::net::TcpListener;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;

use httpmock::Method::{GET, POST};
use httpmock::MockServer;
use linkshieldai::{LinkShieldAI, LinkShieldAIError};

fn environment_lock() -> &'static Mutex<()> {
    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
    LOCK.get_or_init(|| Mutex::new(()))
}

fn client(server: &MockServer) -> LinkShieldAI {
    LinkShieldAI::builder()
        .api_key("test-key")
        .base_url(server.base_url())
        .backoff_factor(Duration::ZERO)
        .build()
        .unwrap()
}

#[test]
fn requires_api_key() {
    let _guard = environment_lock().lock().unwrap();
    std::env::remove_var("LINKSHIELDAI_API_KEY");
    let error = LinkShieldAI::new("").unwrap_err();
    assert!(matches!(error, LinkShieldAIError::Authentication));
}

#[tokio::test]
async fn scan_posts_bearer_authenticated_v1_request() {
    let server = MockServer::start();
    let mock = server.mock(|when, then| {
        when.method(POST).path("/v1/scan").header("authorization", "Bearer test-key");
        then.status(200).json_body_obj(&serde_json::json!({"request_id":"req-1","url":{"submitted":"https://example.com","normalized":"https://example.com/","redirects":[]},"mode":"standard","verdict":"SAFE","confidence":null,"risk_score":null,"threat_categories":[],"reason_codes":["KNOWN_SAFE_MATCH"],"brand_target":null,"screenshot_url":null,"scanned_at":"2026-06-21T00:00:00.000Z","freshness":"live","engine_version":"linkshield-v1"}));
    });
    let result = client(&server)
        .scan("https://example.com", "standard")
        .await
        .unwrap();
    assert_eq!(result.verdict, "SAFE");
    mock.assert();
}

#[tokio::test]
async fn uses_api_key_from_environment() {
    let _guard = environment_lock().lock().unwrap();
    std::env::set_var("LINKSHIELDAI_API_KEY", "env-key");
    let server = MockServer::start();
    let mock = server.mock(|when, then| {
        when.method(GET)
            .path("/")
            .header("authorization", "Bearer env-key")
            .query_param("url", "https://example.com");
        then.status(200)
            .json_body_obj(&serde_json::json!({"result": "likely safe"}));
    });

    let client = LinkShieldAI::builder()
        .api_key("")
        .base_url(server.base_url())
        .build()
        .unwrap();
    client.basic_check("https://example.com").await.unwrap();

    mock.assert();
    std::env::remove_var("LINKSHIELDAI_API_KEY");
}

#[tokio::test]
async fn basic_check_builds_query_and_parses_malicious() {
    let server = MockServer::start();
    let mock = server.mock(|when, then| {
        when.method(GET)
            .path("/")
            .header("authorization", "Bearer test-key")
            .query_param("url", "https://bad.test/login");
        then.status(200)
            .json_body_obj(&serde_json::json!({"result": "Might be malicious"}));
    });

    let result = client(&server)
        .basic_check("https://bad.test/login")
        .await
        .unwrap();

    mock.assert();
    assert!(result.is_malicious);
    assert!(!result.is_safe);
}

#[tokio::test]
async fn basic_check_safe_response() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/");
        then.status(200)
            .json_body_obj(&serde_json::json!({"result": "likely safe"}));
    });

    let result = client(&server)
        .basic_check("https://example.com")
        .await
        .unwrap();

    assert!(!result.is_malicious);
    assert!(result.is_safe);
    assert_eq!(result.raw["result"], "likely safe");
}

#[tokio::test]
async fn detailed_check_normalizes_screenshot_url_and_tag() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/classify_link");
        then.status(200).json_body_obj(&serde_json::json!({
            "result": "Might be malicious",
            "screenshot url": "https://api.linkshieldai.com/screenshot/abc.png",
            "tag": "Yahoo"
        }));
    });

    let result = client(&server)
        .detailed_check("https://phish.test")
        .await
        .unwrap();

    assert!(result.is_malicious);
    assert_eq!(
        result.screenshot_url.as_deref(),
        Some("https://api.linkshieldai.com/screenshot/abc.png")
    );
    assert_eq!(result.tag.as_deref(), Some("Yahoo"));
}

#[tokio::test]
async fn screenshot_downloads_bytes_and_writes_file() {
    let server = MockServer::start();
    let mock = server.mock(|when, then| {
        when.method(GET)
            .path("/screenshot/abc.png")
            .header("authorization", "Bearer test-key");
        then.status(200).body("PNGDATA");
    });
    let output = PathBuf::from(std::env::temp_dir()).join("linkshieldai-rust-shot.png");

    let content = client(&server)
        .get_screenshot(
            "https://api.linkshieldai.com/screenshot/abc.png",
            Some(&output),
        )
        .await
        .unwrap();

    mock.assert();
    assert_eq!(content, b"PNGDATA");
    assert_eq!(std::fs::read(&output).unwrap(), b"PNGDATA");
    let _ = std::fs::remove_file(output);
}

#[tokio::test]
async fn nsfw_check_converts_string_boolean() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/nsfw/site");
        then.status(200)
            .json_body_obj(&serde_json::json!({"result": "True"}));
    });

    let result = client(&server)
        .nsfw_check("https://adult.test")
        .await
        .unwrap();

    assert!(result.is_nsfw);
}

#[tokio::test]
async fn chimera_parses_ai_response() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/chimera");
        then.status(200).json_body_obj(&serde_json::json!({
            "detection_method": "ai-model",
            "matched_signatures": 0,
            "probability": 0.9947241392537983_f64,
            "result": "malicious",
            "url": "https://steamcomun1ty.com/id=5944478616"
        }));
    });

    let result = client(&server)
        .chimera("https://steamcomun1ty.com/id=5944478616")
        .await
        .unwrap();

    assert!(result.is_malicious);
    let probability = result.probability.unwrap();
    assert!((probability - 0.9947241392537983).abs() < 0.000000000000001);
    assert_eq!(result.detection_method.as_deref(), Some("ai-model"));
    assert_eq!(result.matched_signatures, Some(0));
}

#[tokio::test]
async fn error_payload_raises_response_error() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/classify_link");
        then.status(200)
            .json_body_obj(&serde_json::json!({"Error": "Unable to reach the site"}));
    });

    let error = client(&server)
        .detailed_check("https://down.test")
        .await
        .unwrap_err();

    assert!(matches!(error, LinkShieldAIError::Response { .. }));
    assert_eq!(error.to_string(), "Unable to reach the site");
}

#[tokio::test]
async fn rate_limit_raises_rate_limit_error() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/");
        then.status(429)
            .json_body_obj(&serde_json::json!({"error": "Too many requests"}));
    });

    let client = LinkShieldAI::builder()
        .api_key("test-key")
        .base_url(server.base_url())
        .max_retries(0)
        .build()
        .unwrap();
    let error = client.basic_check("https://example.com").await.unwrap_err();

    assert!(matches!(error, LinkShieldAIError::RateLimit { .. }));
}

#[tokio::test]
async fn http_error_raises_status_error() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/chimera");
        then.status(500)
            .json_body_obj(&serde_json::json!({"error": "Invalid or unreachable URL."}));
    });

    let client = LinkShieldAI::builder()
        .api_key("test-key")
        .base_url(server.base_url())
        .max_retries(0)
        .build()
        .unwrap();
    let error = client.chimera("https://bad.test").await.unwrap_err();

    match error {
        LinkShieldAIError::Status { status_code, .. } => assert_eq!(status_code, 500),
        other => panic!("unexpected error: {other:?}"),
    }
}

#[tokio::test]
async fn retries_transient_status_then_succeeds() {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let address = listener.local_addr().unwrap();
    let calls = Arc::new(AtomicUsize::new(0));
    let server_calls = Arc::clone(&calls);
    std::thread::spawn(move || {
        for stream in listener.incoming().take(2) {
            let mut stream = stream.unwrap();
            let mut buffer = [0_u8; 2048];
            let _ = stream.read(&mut buffer);
            let count = server_calls.fetch_add(1, Ordering::SeqCst);
            let (status, body) = if count == 0 {
                ("503 Service Unavailable", r#"{"error":"temporary"}"#)
            } else {
                ("200 OK", r#"{"result":"likely safe"}"#)
            };
            let response = format!(
                "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
                body.len()
            );
            let _ = stream.write_all(response.as_bytes());
        }
    });

    let client = LinkShieldAI::builder()
        .api_key("test-key")
        .base_url(format!("http://{address}"))
        .backoff_factor(Duration::ZERO)
        .build()
        .unwrap();

    let result = client.basic_check("https://example.com").await.unwrap();

    assert!(result.is_safe);
    assert_eq!(calls.load(Ordering::SeqCst), 2);
}

#[tokio::test]
async fn invalid_html_response_includes_preview() {
    let server = MockServer::start();
    server.mock(|when, then| {
        when.method(GET).path("/");
        then.status(200)
            .header("content-type", "text/html")
            .body("<html>oops</html>");
    });

    let error = client(&server)
        .basic_check("https://example.com")
        .await
        .unwrap_err();

    assert!(matches!(error, LinkShieldAIError::Response { .. }));
    assert!(error.to_string().contains("Body preview"));
}