musq 0.0.4

use std::{
    collections::HashSet,
    fmt,
    ops::{Deref, DerefMut},
};

use async_trait::async_trait;
use either::Either;
use futures_core::stream::BoxStream;

use crate::{
    Arguments, QueryResult, Result, Row,
    encode::Encode,
    executor::Execute,
    pool::PoolConnection,
    sqlite::{Value, statement::Statement},
};

/// Raw SQL query with bind parameters. Returned by [`query`].
#[must_use = "query must be executed to affect database"]
#[derive(Clone)]
pub struct Query {
    /// SQL text or prepared statement reference.
    pub(crate) statement: Either<String, Statement>,
    /// Bound arguments for the query.
    pub(crate) arguments: Option<Arguments>,
    /// Whether the query contains raw SQL fragments.
    pub(crate) tainted: bool,
}

impl fmt::Debug for Query {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write_query_debug(f, self.sql(), &self.arguments, self.tainted)
    }
}

/// Render the debug representation for a query.
fn write_query_debug(
    f: &mut fmt::Formatter<'_>,
    sql: &str,
    arguments: &Option<Arguments>,
    tainted: bool,
) -> fmt::Result {
    let multi_line = sql.contains('\n');

    if multi_line && f.alternate() {
        // Pretty multi-line format
        writeln!(f, "Query {{")?;
        writeln!(f, "    statement:")?;
        for line in sql.lines() {
            writeln!(f, "        {line}")?;
        }
        write_optional_fields(f, arguments, tainted, "    ", ",\n")?;
        write!(f, "}}")
    } else if multi_line {
        // Compact multi-line format
        write!(f, "Query {{ statement: ")?;
        for (i, line) in sql.lines().enumerate() {
            if i == 0 {
                writeln!(f, "{line}")?
            } else {
                writeln!(f, "    {line}")?
            }
        }
        write_optional_fields(f, arguments, tainted, ", ", " ")?;
        write!(f, "}}")
    } else {
        // Single-line format
        let mut debug_struct = f.debug_struct("Query");
        debug_struct.field("statement", &sql);
        if let Some(args) = arguments
            && !args.values.is_empty()
        {
            debug_struct.field("arguments", &format_args!("{}", FormatArguments(args)));
        }
        if tainted {
            debug_struct.field("tainted", &tainted);
        }
        debug_struct.finish()
    }
}

/// Write optional argument and taint fields to the formatter.
fn write_optional_fields(
    f: &mut fmt::Formatter<'_>,
    arguments: &Option<Arguments>,
    tainted: bool,
    prefix: &str,
    suffix: &str,
) -> fmt::Result {
    if let Some(args) = arguments
        && !args.values.is_empty()
    {
        write!(
            f,
            "{}arguments: {}{}",
            prefix,
            FormatArguments(args),
            suffix
        )?;
    }
    if tainted {
        write!(f, "{prefix}tainted: {tainted}{suffix}")?;
    }
    Ok(())
}

/// Helper for formatting arguments in debug output.
struct FormatArguments<'a>(&'a Arguments);

impl<'a> fmt::Display for FormatArguments<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let named_indices: HashSet<_> = self.0.named.values().copied().collect();
        let mut first = true;

        write!(f, "[")?;

        // Positional parameters (excluding named ones)
        for (i, value) in self.0.values.iter().enumerate() {
            if !named_indices.contains(&(i + 1)) {
                if !first {
                    write!(f, ", ")?
                }
                write!(f, "{}", FormatValue(value))?;
                first = false;
            }
        }

        // Named parameters
        let mut named: Vec<_> = self.0.named.iter().collect();
        named.sort_by_key(|&(_, &idx)| idx);

        for (name, &idx) in named {
            if !first {
                write!(f, ", ")?
            }
            if let Some(value) = self.0.values.get(idx - 1) {
                write!(f, "{}={}", name, FormatValue(value))?;
                first = false;
            }
        }

        write!(f, "]")
    }
}

/// Helper for formatting SQLite values in debug output.
struct FormatValue<'a>(&'a Value);

impl<'a> fmt::Display for FormatValue<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.0 {
            Value::Null { .. } => write!(f, "NULL"),
            Value::Integer { value, .. } => write!(f, "{value}"),
            Value::Double { value, .. } => write!(f, "{value}"),
            Value::Text { .. } => match self.0.text() {
                Ok(text) => write!(f, "{text:?}"),
                Err(err) => write!(f, "<invalid utf-8: {err}>"),
            },
            Value::Blob { value, .. } => {
                write!(f, "0x")?;
                let display_bytes =
                    value
                        .iter()
                        .take(if value.len() <= 16 { value.len() } else { 8 });
                for byte in display_bytes {
                    write!(f, "{byte:02x}")?;
                }
                if value.len() > 16 {
                    write!(f, "...({} bytes)", value.len())?;
                }
                Ok(())
            }
        }
    }
}

/// SQL query that will map its results to owned Rust types.
///
/// Returned by [`Query::try_map`], `query!()`, etc. Has most of the same methods as [`Query`] but
/// the return types are changed to reflect the mapping. However, there is no equivalent of
/// [`Query::execute`] as it doesn't make sense to map the result type and then ignore it.
///
/// [`Map::bind`] and [`Map::bind_named`] may be used to add parameters after
/// [`Map::try_map`]. Stylistically we still recommend placing your `.bind()` calls
/// before `.try_map()` to avoid adding superfluous binds when using
/// `query!()` et al.
#[must_use = "query must be executed to affect database"]
pub struct Map<F> {
    /// Underlying query.
    inner: Query,
    /// Row mapper function.
    mapper: F,
}

/// Execute queries without exposing the legacy `Executor` trait.
#[async_trait]
pub trait QueryExecutor {
    /// Execute the query and return a summary of changes.
    async fn execute_query(self, query: Query) -> Result<QueryResult>;
    /// Execute the query and stream rows as they are produced.
    fn fetch_query<'c>(self, query: Query) -> BoxStream<'c, Result<Row>>
    where
        Self: 'c;
    /// Execute the query and collect all rows.
    async fn fetch_all_query(self, query: Query) -> Result<Vec<Row>>;
    /// Execute the query and fetch exactly one row.
    async fn fetch_one_query(self, query: Query) -> Result<Row>;
    /// Execute the query and fetch at most one row.
    async fn fetch_optional_query(self, query: Query) -> Result<Option<Row>>;
}

// Implement QueryExecutor for &Pool
#[async_trait]
impl QueryExecutor for &crate::Pool {
    async fn execute_query(self, query: Query) -> Result<QueryResult> {
        let conn = self.acquire().await?;
        conn.execute(query).await
    }

    fn fetch_query<'c>(self, query: Query) -> BoxStream<'c, Result<Row>>
    where
        Self: 'c,
    {
        use futures_util::TryStreamExt;
        Box::pin(async_stream::try_stream! {
            let conn = self.acquire().await?;
            let mut stream = conn.fetch(query);
            while let Some(row) = stream.try_next().await? {
                yield row;
            }
        })
    }

    async fn fetch_all_query(self, query: Query) -> Result<Vec<Row>> {
        let conn = self.acquire().await?;
        conn.fetch_all(query).await
    }

    async fn fetch_one_query(self, query: Query) -> Result<Row> {
        let conn = self.acquire().await?;
        conn.fetch_one(query).await
    }

    async fn fetch_optional_query(self, query: Query) -> Result<Option<Row>> {
        let conn = self.acquire().await?;
        conn.fetch_optional(query).await
    }
}

// Implement QueryExecutor for &Connection
#[async_trait]
impl QueryExecutor for &crate::Connection {
    async fn execute_query(self, query: Query) -> Result<QueryResult> {
        self.execute(query).await
    }

    fn fetch_query<'c>(self, query: Query) -> BoxStream<'c, Result<Row>>
    where
        Self: 'c,
    {
        self.fetch(query)
    }

    async fn fetch_all_query(self, query: Query) -> Result<Vec<Row>> {
        self.fetch_all(query).await
    }

    async fn fetch_one_query(self, query: Query) -> Result<Row> {
        self.fetch_one(query).await
    }

    async fn fetch_optional_query(self, query: Query) -> Result<Option<Row>> {
        self.fetch_optional(query).await
    }
}

// Implement QueryExecutor for &PoolConnection
#[async_trait]
impl QueryExecutor for &PoolConnection {
    async fn execute_query(self, query: Query) -> Result<QueryResult> {
        self.execute(query).await
    }

    fn fetch_query<'c>(self, query: Query) -> BoxStream<'c, Result<Row>>
    where
        Self: 'c,
    {
        self.fetch(query)
    }

    async fn fetch_all_query(self, query: Query) -> Result<Vec<Row>> {
        self.fetch_all(query).await
    }

    async fn fetch_one_query(self, query: Query) -> Result<Row> {
        self.fetch_one(query).await
    }

    async fn fetch_optional_query(self, query: Query) -> Result<Option<Row>> {
        self.fetch_optional(query).await
    }
}

// Implement QueryExecutor for &Transaction<C>
#[async_trait]
impl<C> QueryExecutor for &crate::Transaction<C>
where
    C: DerefMut<Target = crate::Connection> + Send + Sync,
{
    async fn execute_query(self, query: Query) -> Result<QueryResult> {
        let conn: &crate::Connection = self.deref();
        conn.execute(query).await
    }

    fn fetch_query<'c>(self, query: Query) -> BoxStream<'c, Result<Row>>
    where
        Self: 'c,
    {
        use futures_util::TryStreamExt;
        Box::pin(async_stream::try_stream! {
            let conn: &crate::Connection = self.deref();
            let mut stream = conn.fetch(query);
            while let Some(row) = stream.try_next().await? {
                yield row;
            }
        })
    }

    async fn fetch_all_query(self, query: Query) -> Result<Vec<Row>> {
        let conn: &crate::Connection = self.deref();
        conn.fetch_all(query).await
    }

    async fn fetch_one_query(self, query: Query) -> Result<Row> {
        let conn: &crate::Connection = self.deref();
        conn.fetch_one(query).await
    }

    async fn fetch_optional_query(self, query: Query) -> Result<Option<Row>> {
        let conn: &crate::Connection = self.deref();
        conn.fetch_optional(query).await
    }
}

// Implement QueryExecutor for &mut Transaction<C>
#[async_trait]
impl<C> QueryExecutor for &mut crate::Transaction<C>
where
    C: DerefMut<Target = crate::Connection> + Send + Sync,
{
    async fn execute_query(self, query: Query) -> Result<QueryResult> {
        let conn: &crate::Connection = self.deref();
        conn.execute(query).await
    }

    fn fetch_query<'c>(self, query: Query) -> BoxStream<'c, Result<Row>>
    where
        Self: 'c,
    {
        use futures_util::TryStreamExt;
        Box::pin(async_stream::try_stream! {
            let conn: &crate::Connection = self.deref();
            let mut stream = conn.fetch(query);
            while let Some(row) = stream.try_next().await? {
                yield row;
            }
        })
    }

    async fn fetch_all_query(self, query: Query) -> Result<Vec<Row>> {
        let conn: &crate::Connection = self.deref();
        conn.fetch_all(query).await
    }

    async fn fetch_one_query(self, query: Query) -> Result<Row> {
        let conn: &crate::Connection = self.deref();
        conn.fetch_one(query).await
    }

    async fn fetch_optional_query(self, query: Query) -> Result<Option<Row>> {
        let conn: &crate::Connection = self.deref();
        conn.fetch_optional(query).await
    }
}

impl Execute for Query {
    fn sql(&self) -> &str {
        match &self.statement {
            Either::Right(statement) => statement.sql(),
            Either::Left(sql) => sql,
        }
    }

    fn arguments(&mut self) -> Option<Arguments> {
        self.arguments.take()
    }
}

impl<F> Map<F> {
    /// Attempt to bind a value for use with the mapped query.
    pub fn try_bind<'q, T: 'q + Send + Encode>(mut self, value: T) -> Result<Self> {
        self.inner = self.inner.try_bind(value)?;
        Ok(self)
    }

    /// Bind a value for use with the mapped query.
    ///
    /// This will panic if [`try_bind`](Self::try_bind) returns an error.
    pub fn bind<'q, T: 'q + Send + Encode>(self, value: T) -> Self {
        self.try_bind(value)
            .expect("failed to bind query parameter")
    }

    /// Attempt to bind a value to a named parameter.
    pub fn try_bind_named<'q, T: 'q + Send + Encode>(
        mut self,
        name: &str,
        value: T,
    ) -> Result<Self> {
        self.inner = self.inner.try_bind_named(name, value)?;
        Ok(self)
    }

    /// Bind a value to a named parameter.
    ///
    /// This will panic if [`try_bind_named`](Self::try_bind_named) returns an error.
    pub fn bind_named<'q, T: 'q + Send + Encode>(self, name: &str, value: T) -> Self {
        self.try_bind_named(name, value)
            .expect("failed to bind named query parameter")
    }
}

impl Query {
    /// Returns `true` if the query has had raw SQL appended to it.
    pub fn is_tainted(&self) -> bool {
        self.tainted
    }

    /// Convert this query into a mutable [`crate::QueryBuilder`].
    pub fn into_builder(self) -> crate::QueryBuilder {
        crate::QueryBuilder::from_parts(
            self.sql().to_string(),
            self.arguments.unwrap_or_default(),
            self.tainted,
        )
    }

    /// Joins this query with another [`Query`].
    ///
    /// The SQL and arguments from `other` are appended to this query and a new
    /// combined query is returned.
    ///
    /// This method panics if the appended query contains numeric positional
    /// placeholders such as `?1` or numeric `$1`. Use [`Query::try_join`] to
    /// handle unsupported composition as an error.
    pub fn join(self, other: Self) -> Self {
        self.try_join(other)
            .expect("failed to join query fragments")
    }

    /// Try to join this query with another [`Query`].
    ///
    /// Anonymous `?` and named parameters are supported in composed fragments.
    /// Numeric positional placeholders such as `?1` and numeric `$1` are
    /// rejected in appended fragments because their SQLite indices are absolute
    /// within the final statement and cannot be safely rebased.
    pub fn try_join(self, other: Self) -> Result<Self> {
        let mut builder = self.into_builder();
        builder.try_push_query(other)?;
        Ok(builder.build())
    }

    /// Attempt to bind a value for use with this SQL query.
    ///
    /// If the number of times this is called does not match the number of bind parameters that
    /// appear in the query then an error will be returned when this query is executed.
    pub fn try_bind<'q, T: 'q + Send + Encode>(mut self, value: T) -> Result<Self> {
        if let Some(arguments) = &mut self.arguments {
            arguments.add(&value)?;
        }
        drop(value);
        Ok(self)
    }

    /// Bind a value for use with this SQL query.
    ///
    /// This will panic if [`try_bind`](Self::try_bind) returns an error.
    pub fn bind<'q, T: 'q + Send + Encode>(self, value: T) -> Self {
        self.try_bind(value)
            .expect("failed to bind query parameter")
    }

    /// Attempt to bind a value to a named parameter.
    pub fn try_bind_named<'q, T: 'q + Send + Encode>(
        mut self,
        name: &str,
        value: T,
    ) -> Result<Self> {
        if let Some(arguments) = &mut self.arguments {
            arguments.add_named(name, &value)?;
        }
        drop(value);
        Ok(self)
    }

    /// Bind a value to a named parameter.
    ///
    /// This will panic if [`try_bind_named`](Self::try_bind_named) returns an error.
    pub fn bind_named<'q, T: 'q + Send + Encode>(self, name: &str, value: T) -> Self {
        self.try_bind_named(name, value)
            .expect("failed to bind named query parameter")
    }

    /// Map each row in the result to another type.
    ///
    /// See [`try_map`](Query::try_map) for a fallible version of this method.
    ///
    /// The [`query_as`] function will construct a mapped query using
    /// a [`FromRow`] implementation.
    pub fn map<F, O>(self, mut f: F) -> Map<impl FnMut(Row) -> Result<O> + Send>
    where
        F: FnMut(Row) -> O + Send,
        O: Unpin,
    {
        self.try_map(move |row| Ok(f(row)))
    }

    /// Map each row in the result to another type.
    ///
    /// The [`query_as`] function will construct a mapped query using
    /// a [`FromRow`] implementation.
    pub fn try_map<F, O>(self, f: F) -> Map<F>
    where
        F: FnMut(Row) -> Result<O> + Send,
        O: Unpin,
    {
        Map {
            inner: self,
            mapper: f,
        }
    }

    /// Execute the query and return the total number of rows affected.
    pub async fn execute<E>(self, executor: E) -> Result<QueryResult>
    where
        E: QueryExecutor,
    {
        executor.execute_query(self).await
    }

    /// Execute the query and return the generated results as a stream.
    pub fn fetch<'c, E>(self, executor: E) -> BoxStream<'c, Result<Row>>
    where
        E: QueryExecutor + 'c,
    {
        executor.fetch_query(self)
    }

    /// Execute the query and return all the generated results, collected into a [`Vec`].
    pub async fn fetch_all<E>(self, executor: E) -> Result<Vec<Row>>
    where
        E: QueryExecutor,
    {
        executor.fetch_all_query(self).await
    }

    /// Execute the query and returns exactly one row.
    pub async fn fetch_one<E>(self, executor: E) -> Result<Row>
    where
        E: QueryExecutor,
    {
        executor.fetch_one_query(self).await
    }

    /// Execute the query and returns at most one row.
    pub async fn fetch_optional<E>(self, executor: E) -> Result<Option<Row>>
    where
        E: QueryExecutor,
    {
        executor.fetch_optional_query(self).await
    }
}

impl<F: Send> Execute for Map<F> {
    fn sql(&self) -> &str {
        self.inner.sql()
    }

    fn arguments(&mut self) -> Option<Arguments> {
        self.inner.arguments()
    }
}

impl<F, O> Map<F>
where
    F: FnMut(Row) -> Result<O> + Send,
    O: Send + Unpin,
{
    /// Map each row in the result to another type.
    ///
    /// See [`try_map`](Map::try_map) for a fallible version of this method.
    ///
    /// The [`query_as`] function will construct a mapped query using
    /// a [`FromRow`] implementation.
    pub fn map<G, P>(self, mut g: G) -> Map<impl FnMut(Row) -> Result<P> + Send>
    where
        G: FnMut(O) -> P + Send,
        P: Unpin,
    {
        self.try_map(move |data| Ok(g(data)))
    }

    /// Map each row in the result to another type.
    ///
    /// The [`query_as`] function will construct a mapped query using
    /// a [`FromRow`] implementation.
    pub fn try_map<G, P>(self, mut g: G) -> Map<impl FnMut(Row) -> Result<P> + Send>
    where
        G: FnMut(O) -> Result<P> + Send,
        P: Unpin,
    {
        let mut f = self.mapper;
        Map {
            inner: self.inner,
            mapper: move |row| f(row).and_then(&mut g),
        }
    }

    /// Execute the query and return all the generated results, collected into a [`Vec`].
    pub async fn fetch_all<E>(mut self, executor: E) -> Result<Vec<O>>
    where
        E: QueryExecutor,
    {
        let rows = self.inner.fetch_all(executor).await?;
        let mut results = Vec::with_capacity(rows.len());
        for row in rows {
            results.push((self.mapper)(row)?);
        }
        Ok(results)
    }

    /// Execute the query and returns exactly one row.
    pub async fn fetch_one<E>(self, executor: E) -> Result<O>
    where
        E: QueryExecutor,
    {
        match self.fetch_optional(executor).await? {
            Some(row) => Ok(row),
            None => Err(crate::Error::RowNotFound),
        }
    }

    /// Execute the query and returns at most one row.
    pub async fn fetch_optional<E>(mut self, executor: E) -> Result<Option<O>>
    where
        E: QueryExecutor,
    {
        let row = self.inner.fetch_optional(executor).await?;

        if let Some(row) = row {
            (self.mapper)(row).map(Some)
        } else {
            Ok(None)
        }
    }
}

/// Make a SQL query from a prepared statement.
pub(crate) fn query_statement(statement: &Statement) -> Query {
    Query {
        arguments: Some(Default::default()),
        statement: Either::Right(statement.clone()),
        tainted: false,
    }
}

/// Make a SQL query from a prepared statement with the given arguments.
pub(crate) fn query_statement_with(statement: &Statement, arguments: Arguments) -> Query {
    Query {
        arguments: Some(arguments),
        statement: Either::Right(statement.clone()),
        tainted: false,
    }
}

/// Make a SQL query.
pub fn query(sql: &str) -> Query {
    Query {
        arguments: Some(Default::default()),
        statement: Either::Left(sql.to_string()),
        tainted: false,
    }
}

/// Make a SQL query, with the given arguments.
pub fn query_with(sql: &str, arguments: Arguments) -> Query {
    Query {
        arguments: Some(arguments),
        statement: Either::Left(sql.to_string()),
        tainted: false,
    }
}

use crate::from_row::FromRow;

/// Build a typed query that maps rows into `O` via [`FromRow`].
pub fn query_as<'q, O>(sql: &'q str) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    O: Send + Unpin + for<'r> FromRow<'r>,
{
    query(sql).try_map(|row| O::from_row("", &row))
}

/// Build a typed query with explicit arguments.
pub fn query_as_with<'q, O>(
    sql: &'q str,
    arguments: Arguments,
) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    O: Send + Unpin + for<'r> FromRow<'r>,
{
    query_with(sql, arguments).try_map(|row| O::from_row("", &row))
}

/// Build a typed query from a prepared statement.
pub(crate) fn query_statement_as<'q, O>(
    statement: &'q Statement,
) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    O: Send + Unpin + for<'r> FromRow<'r>,
{
    query_statement(statement).try_map(|row| O::from_row("", &row))
}

/// Build a typed query from a prepared statement and arguments.
pub(crate) fn query_statement_as_with<'q, O>(
    statement: &'q Statement,
    arguments: Arguments,
) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    O: Send + Unpin + for<'r> FromRow<'r>,
{
    query_statement_with(statement, arguments).try_map(|row| O::from_row("", &row))
}

/// Build a scalar query that maps a single column into `O`.
pub fn query_scalar<'q, O>(sql: &'q str) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    (O,): for<'r> FromRow<'r>,
    O: Send + Unpin,
{
    query_as(sql).map(|(o,)| o)
}

/// Build a scalar query with explicit arguments.
pub fn query_scalar_with<'q, O>(
    sql: &'q str,
    arguments: Arguments,
) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    (O,): for<'r> FromRow<'r>,
    O: Send + Unpin,
{
    query_as_with(sql, arguments).map(|(o,)| o)
}

/// Build a scalar query from a prepared statement.
pub(crate) fn query_statement_scalar<'q, O>(
    statement: &'q Statement,
) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    (O,): for<'r> FromRow<'r>,
    O: Send + Unpin,
{
    query_statement_as(statement).map(|(o,)| o)
}

/// Build a scalar query from a prepared statement and arguments.
pub(crate) fn query_statement_scalar_with<'q, O>(
    statement: &'q Statement,
    arguments: Arguments,
) -> Map<impl FnMut(Row) -> Result<O> + Send>
where
    (O,): for<'r> FromRow<'r>,
    O: Send + Unpin,
{
    query_statement_as_with(statement, arguments).map(|(o,)| o)
}

/// Quote an identifier for use in a SQL statement.
pub fn quote_identifier(ident: &str) -> String {
    let escaped = ident.replace('"', "\"\"");
    format!("\"{escaped}\"")
}