better_sparql_client 0.2.1

use oxrdf::Graph;
use oxttl::TurtleParseError;
use sparesults::{
    QueryResultsFormat, QueryResultsParseError, QueryResultsParser, QuerySolution,
    ReaderQueryResultsParserOutput,
};
use url::Url;

use crate::utils::parse_ttl_into_graph;

/// Errors that can occur related to transport (e.g. network issues, unexpected status codes, etc.)
#[derive(Debug, thiserror::Error)]
pub enum TransportError {
    #[error("HTTP error: {0}")]
    Http(#[from] reqwest::Error),

    #[error("Received unexpected status code: {0}")]
    BadStatusCode(reqwest::StatusCode),
}

#[derive(Debug, thiserror::Error)]
pub enum ResponseContentTypeError {
    #[error("Missing content type in response")]
    MissingContentType,

    #[error("Unexpected content type in response: {0}")]
    UnexpectedContentType(String),
}

#[derive(Debug, thiserror::Error)]
pub enum QueryError {
    #[error(transparent)]
    Transport(#[from] TransportError),

    #[error(transparent)]
    ResponseContentType(#[from] ResponseContentTypeError),

    #[error("Failed to parse results: {0}")]
    FailedToParse(#[from] QueryResultsParseError),

    #[error("Unexpected query results format. Expected query solutions, got boolean.")]
    UnexpectedFormat,
}

#[derive(Debug, thiserror::Error)]
pub enum QueryGraphError {
    #[error(transparent)]
    Transport(#[from] TransportError),

    #[error(transparent)]
    ResponseContentType(#[from] ResponseContentTypeError),

    #[error("Failed to parse graph: {0}")]
    FailedToParse(#[from] TurtleParseError),
}

#[derive(Debug, thiserror::Error)]
pub enum AskError {
    #[error(transparent)]
    Transport(#[from] TransportError),

    #[error(transparent)]
    ResponseContentType(#[from] ResponseContentTypeError),

    #[error("Failed to parse result: {0}")]
    FailedToParse(#[from] QueryResultsParseError),

    #[error("Unexpected query results format. Expected boolean, got query solutions.")]
    UnexpectedFormat,
}

// update could probably just return a TransportError, but this is consistent and leaves room for future error varints if ever necessary.
#[derive(Debug, thiserror::Error)]
pub enum UpdateError {
    #[error(transparent)]
    Transport(#[from] TransportError),
}

/// validates that a response has a successful status code. otherwise, returns a transport error
fn validate_response_status(response: &reqwest::Response) -> Result<(), TransportError> {
    if !response.status().is_success() {
        return Err(TransportError::BadStatusCode(response.status()));
    }

    Ok(())
}

/// validates that a response has the expected content type. otherwise, returns a transport error
fn validate_response_content_type(
    response: &reqwest::Response,
    expected_content_type: &str,
) -> Result<(), ResponseContentTypeError> {
    let Some(content_type) = response.headers().get("content-type") else {
        return Err(ResponseContentTypeError::MissingContentType);
    };

    let content_type_str = content_type.to_str().unwrap_or("");

    if !content_type_str.contains(expected_content_type) {
        return Err(ResponseContentTypeError::UnexpectedContentType(
            content_type_str.to_owned(),
        ));
    }

    Ok(())
}

#[derive(Debug, Clone)]
pub struct SparqlClient {
    http_client: reqwest::Client,
    service_url: Url,
}

impl SparqlClient {
    pub fn new(service_url: Url) -> Self {
        Self {
            http_client: reqwest::Client::new(),
            service_url,
        }
    }

    pub fn with_http_client(service_url: Url, http_client: reqwest::Client) -> Self {
        Self {
            http_client,
            service_url,
        }
    }

    fn build_query_request(&self, query: String) -> reqwest::RequestBuilder {
        self.http_client
            .post(self.service_url.clone())
            .header("Content-Type", "application/sparql-query")
            .body(query)
    }

    /// Executes a SELECT query and returns the results as a vector of QuerySolution.
    pub async fn query(&self, query: &str) -> Result<Vec<QuerySolution>, QueryError> {
        let response = self
            .build_query_request(query.to_owned())
            .header("Accept", "application/sparql-results+json")
            .send()
            .await
            .map_err(TransportError::Http)?;

        validate_response_status(&response)?;
        validate_response_content_type(&response, "application/sparql-results+json")?;

        let body_bytes = response.bytes().await.map_err(TransportError::Http)?;
        let query_results_parser = QueryResultsParser::from_format(QueryResultsFormat::Json);

        let ReaderQueryResultsParserOutput::Solutions(solutions) =
            query_results_parser.for_reader(&*body_bytes)?
        else {
            return Err(QueryError::UnexpectedFormat);
        };

        let solutions = solutions.into_iter().flatten().collect();

        Ok(solutions)
    }

    /// Executes a CONSTRUCT or DESCRIBE query and returns the resulting graph.
    pub async fn query_graph(&self, query: &str) -> Result<Graph, QueryGraphError> {
        let response = self
            .build_query_request(query.to_owned())
            .header("Accept", "text/turtle")
            .send()
            .await
            .map_err(TransportError::Http)?;

        validate_response_status(&response)?;
        validate_response_content_type(&response, "text/turtle")?;

        let body = response.text().await.map_err(TransportError::Http)?;
        let graph = parse_ttl_into_graph(&body)?;

        Ok(graph)
    }

    /// Executes an ASK query and returns the boolean result.
    pub async fn ask(&self, query: &str) -> Result<bool, AskError> {
        let response = self
            .build_query_request(query.to_owned())
            .header("Accept", "application/sparql-results+json")
            .send()
            .await
            .map_err(TransportError::Http)?;

        validate_response_status(&response)?;
        validate_response_content_type(&response, "application/sparql-results+json")?;

        let body_bytes = response.bytes().await.map_err(TransportError::Http)?;
        let query_results_parser = QueryResultsParser::from_format(QueryResultsFormat::Json);

        let ReaderQueryResultsParserOutput::Boolean(is_true) =
            query_results_parser.for_reader(&*body_bytes)?
        else {
            return Err(AskError::UnexpectedFormat);
        };

        Ok(is_true)
    }

    /// Executes an UPDATE query (INSERT, DELETE, etc.)
    /// Returns a result indicating success or failure of the update operation.
    pub async fn update(&self, query: &str) -> Result<(), UpdateError> {
        let response = self
            .build_query_request(query.to_owned())
            .send()
            .await
            .map_err(TransportError::Http)?;

        // since we are't expecting any particular content type in the response, we only validate the status code here, not the content type.
        validate_response_status(&response)?;

        Ok(())
    }
}