/*!

[Aggregate](https://www.postgresql.org/docs/current/xaggr.html) support.

Most items of this trait map directly to a [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
functionality.

Aggregates are created by implementing [`Aggregate`] for a type and decorating the implementation with
[`#[pg_aggregate]`](pgx_macros::pg_aggregate).

Definition of the aggregate is done via settings in the type's [`Aggregate`] implementation. While
the trait itself has several items, only a few are required, the macro will fill in the others with unused stubs.

# Minimal Example

```rust
use pgx::*;
use serde::{Serialize, Deserialize};

// pg_module_magic!(); // Uncomment this outside of docs!

#[derive(Copy, Clone, Default, PostgresType, Serialize, Deserialize)]
pub struct DemoSum {
    count: i32,
}

#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = i32;
    fn state(
        mut current: Self::State,
        arg: Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        current.count += arg;
        current
    }
}
```

This creates SQL like so:

```sql
-- src/lib.rs:11
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
	integer /* i32 */
)
(
	SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
	STYPE = DemoSum, /* aggregate::DemoSum */
	INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);
```

Example of usage:

```sql
aggregate=# CREATE TABLE demo_table (value INTEGER);
CREATE TABLE
aggregate=# INSERT INTO demo_table (value) VALUES (1), (2), (3);
INSERT 0 3
aggregate=# SELECT DemoSum(value) FROM demo_table;
    demosum    
-------------
 {"count":6}
(1 row)
```

## Multiple Arguments

Sometimes aggregates need to handle multiple arguments. The
[`Aggregate::Args`](Aggregate::Args) associated type can be a tuple:

```rust
# use pgx::*;
# use serde::{Serialize, Deserialize};
# 
# #[derive(Copy, Clone, Default, PostgresType, Serialize, Deserialize)]
# pub struct DemoSum {
#     count: i32,
# }
#
#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = (i32, i32);
    fn state(
        mut current: Self::State,
        (arg1, arg2): Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        current.count += arg1;
        current.count += arg2;
        current
    }
}
```

Creates:

```sql
-- src/lib.rs:11
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
	integer, /* i32 */
	integer /* i32 */
)
(
	SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
	STYPE = DemoSum, /* aggregate::DemoSum */
	INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);
```

## Named Arguments

The [`name!(ident, Type)`][macro@crate::name] macro can be used to set the name of an argument:

```rust
# use pgx::*;
# use serde::{Serialize, Deserialize};
# 
# #[derive(Copy, Clone, Default, PostgresType, Serialize, Deserialize)]
# pub struct DemoSum {
#     count: i32,
# }
# 
# #[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = (
        i32,
        name!(extra, i32),
    );
    fn state(
        mut current: Self::State,
        (arg1, extra): Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        todo!()
    }
}
```

Creates:

```sql
-- src/lib.rs:11
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
	integer, /* i32 */
	"extra" integer /* i32 */
)
(
	SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
	STYPE = DemoSum, /* aggregate::DemoSum */
	INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);
```

## Function attributes

Functions inside the `impl` may use the [`#[pgx]`](macro@crate::pgx) attribute. It
accepts the same parameters as [`#[pg_extern]`][macro@pgx-macros::pg_extern].

```rust
# use pgx::*;
# use serde::{Serialize, Deserialize};
# 
# #[derive(Copy, Clone, Default, PostgresType, Serialize, Deserialize)]
# pub struct DemoSum {
#     count: i32,
# }
#
#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = i32;
    #[pgx(parallel_safe, immutable)]
    fn state(
        mut current: Self::State,
        arg: Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        todo!()
    }
}
```

Generates:

```sql
-- src/lib.rs:11
-- aggregate::demo_sum_state
CREATE OR REPLACE FUNCTION "demo_sum_state"(
	"this" DemoSum, /* aggregate::DemoSum */
	"arg_one" integer /* i32 */
) RETURNS DemoSum /* aggregate::DemoSum */
PARALLEL SAFE IMMUTABLE STRICT
LANGUAGE c /* Rust */
AS 'MODULE_PATHNAME', 'demo_sum_state_wrapper';
```

## Non-`Self` State

Sometimes it's useful to have aggregates share state, or use some other type for state.

```rust
# use pgx::*;
# use serde::{Serialize, Deserialize};
# 
#[derive(Copy, Clone, Default, PostgresType, Serialize, Deserialize)]
pub struct DemoSumState {
    count: i32,
}

pub struct DemoSum;

#[pg_aggregate]
impl Aggregate for DemoSum {
    const INITIAL_CONDITION: Option<&'static str> = Some(r#"{ "count": 0 }"#);
    type Args = i32;
    type State = DemoSumState;
    fn state(
        mut current: Self::State,
        arg: Self::Args,
        _fcinfo: pg_sys::FunctionCallInfo
    ) -> Self::State {
        todo!()
    }
}
```

Creates:

```sql
-- src/lib.rs:13
-- aggregate::demo_sum_state
CREATE OR REPLACE FUNCTION "demo_sum_state"(
	"this" DemoSumState, /* aggregate::DemoSumState */
	"arg_one" integer /* i32 */
) RETURNS DemoSumState /* aggregate::DemoSumState */
STRICT
LANGUAGE c /* Rust */
AS 'MODULE_PATHNAME', 'demo_sum_state_wrapper';

-- src/lib.rs:13
-- aggregate::DemoSum
CREATE AGGREGATE DemoSum (
	integer /* i32 */
)
(
	SFUNC = "demo_sum_state", /* aggregate::DemoSum::state */
	STYPE = DemoSumState, /* aggregate::DemoSumState */
	INITCOND = '{ "count": 0 }' /* aggregate::DemoSum::INITIAL_CONDITION */
);
```

*/

use crate::{
    error,
    datum::sql_entity_graph::{PgxSql, ToSql},
    pgbox::PgBox,
    memcxt::PgMemoryContexts,
    pg_sys::{CurrentMemoryContext, MemoryContext, AggCheckCallContext, FunctionCallInfo},
};

/// Aggregate implementation trait.
/// 
/// When decorated with [`#[pgx_macros::pg_aggregate]`](pgx_macros::pg_aggregate), enables the
/// generation of [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html)
/// SQL.
/// 
/// The [`#[pgx_macros::pg_aggregate]`](pgx_macros::pg_aggregate) will automatically fill fields
/// marked optional with stubs.
pub trait Aggregate
where
    Self: Sized,
{
    /// The type of the return value on `state` and `combine` functions.
    ///
    /// For an aggregate type which does not have a `PgVarlenaInOutFuncs` implementation,
    /// this can be left out, or set to it's default, `Self`.
    ///
    /// For an aggregate type which **does** have a `PgVarlenaInOutFuncs` implementation,
    /// this should be set to `PgVarlena<Self>`.
    ///
    /// Other types are supported as well, this can be useful if multiple aggregates share a state.
    type State;

    /// The type of the argument(s).
    ///
    /// For a single argument, provide the type directly.
    ///
    /// For multiple arguments, provide a tuple.
    ///
    /// Use [`pgx::name!()`](crate::name) to set the SQL name of the argument.
    ///
    /// If the final argument is to be variadic, use [`pgx::variadic`](crate::variadic). When used
    /// with [`pgx::name!()`](crate::name), it must be used **inside** the [`pgx::name!()`](crate::name) macro.
    type Args;

    /// The types of the direct argument(s) to an ordered-set aggregate's `finalize`.
    /// 
    /// **Only effective if `ORDERED_SET` is `true`.**
    ///
    /// For a single argument, provide the type directly.
    ///
    /// For multiple arguments, provide a tuple.
    /// 
    /// For no arguments, don't set this, or set it to `()` (the default).
    ///
    /// `pgx` does not support `argname` as it is only used for documentation purposes.
    ///
    /// If the final argument is to be variadic, use `pgx::Variadic`.
    ///
    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    type OrderedSetArgs;

    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    type Finalize;

    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    type MovingState;

    /// The name of the aggregate. (eg. What you'd pass to `SELECT agg(col) FROM tab`.)
    const NAME: &'static str;

    /// Set to true if this is an ordered set aggregate.
    /// 
    /// If set, the `OrderedSetArgs` associated type becomes effective, this allows for
    /// direct arguments to the `finalize` function.
    /// 
    /// See https://www.postgresql.org/docs/current/xaggr.html#XAGGR-ORDERED-SET-AGGREGATES
    /// for more information.
    /// 
    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const ORDERED_SET: bool = false;

    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const PARALLEL: Option<ParallelOption> = None;

    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const FINALIZE_MODIFY: Option<FinalizeModify> = None;

    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const MOVING_FINALIZE_MODIFY: Option<FinalizeModify> = None;

    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const INITIAL_CONDITION: Option<&'static str> = None;

    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const SORT_OPERATOR: Option<&'static str> = None;

    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const MOVING_INITIAL_CONDITION: Option<&'static str> = None;

    /// **Optional:** This const can be skipped, `#[pg_aggregate]` will create a stub.
    const HYPOTHETICAL: bool = false;

    fn state(current: Self::State, v: Self::Args, fcinfo: FunctionCallInfo) -> Self::State;

    /// The `OrderedSetArgs` is `()` unless `ORDERED_SET` is `true` and `OrderedSetArgs` is configured.
    /// 
    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    fn finalize(current: Self::State, direct_args: Self::OrderedSetArgs, fcinfo: FunctionCallInfo) -> Self::Finalize;

    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    fn combine(current: Self::State, _other: Self::State, fcinfo: FunctionCallInfo) -> Self::State;

    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    fn serial(current: Self::State, fcinfo: FunctionCallInfo) -> Vec<u8>;

    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    fn deserial(current: Self::State, _buf: Vec<u8>, _internal: PgBox<Self::State>, fcinfo: FunctionCallInfo) -> PgBox<Self::State>;

    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    fn moving_state(_mstate: Self::MovingState, _v: Self::Args, fcinfo: FunctionCallInfo) -> Self::MovingState;

    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    fn moving_state_inverse(_mstate: Self::MovingState, _v: Self::Args, fcinfo: FunctionCallInfo) -> Self::MovingState;

    /// The `OrderedSetArgs` is `()` unless `ORDERED_SET` is `true` and `OrderedSetArgs` is configured.
    /// 
    /// **Optional:** This function can be skipped, `#[pg_aggregate]` will create a stub.
    fn moving_finalize(_mstate: Self::MovingState, direct_args: Self::OrderedSetArgs, fcinfo: FunctionCallInfo) -> Self::Finalize;

    unsafe fn memory_context(fcinfo: FunctionCallInfo) -> Option<MemoryContext> {
        if fcinfo.is_null() {
            return Some(CurrentMemoryContext)
        }
        let mut memory_context = std::ptr::null_mut();
        let is_aggregate = AggCheckCallContext(fcinfo, &mut memory_context);
        if is_aggregate == 0 {
            None
        } else {
            debug_assert!(!memory_context.is_null());
            Some(memory_context)
        }
    }

    fn in_memory_context<
        R,
        F: FnOnce(&mut PgMemoryContexts) -> R + std::panic::UnwindSafe + std::panic::RefUnwindSafe
    >(fcinfo: FunctionCallInfo, f: F) -> R {
        let aggregate_memory_context = unsafe {
            Self::memory_context(fcinfo)
        }.unwrap_or_else(|| error!("Cannot access Aggregate memory contexts when not an aggregate."));
        PgMemoryContexts::For(aggregate_memory_context).switch_to(f)
    } 
}

/// Corresponds to the `PARALLEL` and `MFINALFUNC_MODIFY` in [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html).
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum ParallelOption {
    Safe,
    Restricted,
    Unsafe,
}

impl ToSql for ParallelOption {
    fn to_sql(&self, _context: &PgxSql) -> eyre::Result<String> {
        let value = match self {
            ParallelOption::Safe => String::from("SAFE"),
            ParallelOption::Restricted => String::from("RESTRICTED"),
            ParallelOption::Unsafe => String::from("UNSAFE"),
        };
        Ok(value)
    }
}

/// Corresponds to the `FINALFUNC_MODIFY` and `MFINALFUNC_MODIFY` in [`CREATE AGGREGATE`](https://www.postgresql.org/docs/current/sql-createaggregate.html).
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum FinalizeModify {
    ReadOnly,
    Shareable,
    ReadWrite,
}

impl ToSql for FinalizeModify {
    fn to_sql(&self, _context: &PgxSql) -> eyre::Result<String> {
        let value = match self {
            FinalizeModify::ReadOnly => String::from("READ_ONLY"),
            FinalizeModify::Shareable => String::from("SHAREABLE"),
            FinalizeModify::ReadWrite => String::from("READ_WRITE"),
        };
        Ok(value)
    }
}