use chrono::{DateTime, Utc};
use sqlx::{postgres::types::Oid, prelude::FromRow, query_as};
use thiserror::Error;
use uuid::Uuid;

use crate::{client::WorkspaceClient, rolnames::ROLE_PREFIX_USER};

#[derive(Clone, Debug, Eq, Hash, FromRow, PartialEq)]
pub struct PgRole {
    /// ID of role
    pub oid: Oid,

    /// Role name
    pub rolname: String,

    /// Role has superuser privileges
    pub rolsuper: bool,

    /// Role automatically inherits privileges of roles it is a member of
    pub rolinherit: bool,

    /// Role can create more roles
    pub rolcreaterole: bool,

    /// Role can create databases
    pub rolcreatedb: bool,

    /// Role can log in. That is, this role can be given as the initial session authorization identifier
    pub rolcanlogin: bool,

    /// Role is a replication role. A replication role can initiate replication connections and create and drop replication slots.
    pub rolreplication: bool,

    /// For roles that can log in, this sets maximum number of concurrent connections this role can make. -1 means no limit.
    pub rolconnlimit: i32,

    /// Password expiry time (only used for password authentication); null if no expiration
    pub rolvaliduntil: Option<DateTime<Utc>>,

    /// Role bypasses every row-level security policy, see Section 5.9 for more information.
    pub rolbypassrls: bool,
}

impl PgRole {
    pub fn with_name_in(names: Vec<String>) -> WithNameInQuery {
        WithNameInQuery { names }
    }

    pub fn with_name_starting_with(prefix: String) -> WithNameStartingWithQuery {
        WithNameStartingWithQuery { prefix }
    }
}

#[derive(Clone, Debug)]
pub struct WithNameInQuery {
    names: Vec<String>,
}

impl WithNameInQuery {
    pub async fn fetch_all(
        &self,
        client: &mut WorkspaceClient,
    ) -> Result<Vec<PgRole>, sqlx::Error> {
        query_as!(
            PgRole,
            r#"
select
    oid as "oid!",
    rolname as "rolname!",
    rolsuper as "rolsuper!",
    rolinherit as "rolinherit!",
    rolcreaterole as "rolcreaterole!",
    rolcreatedb as "rolcreatedb!",
    rolcanlogin as "rolcanlogin!",
    rolreplication as "rolreplication!",
    rolconnlimit as "rolconnlimit!",
    rolvaliduntil,
    rolbypassrls as "rolbypassrls!"
from pg_roles
where rolname = any($1)
"#,
            self.names.as_slice()
        )
        .fetch_all(client.get_conn())
        .await
    }
}

#[derive(Clone, Debug)]
pub struct WithNameStartingWithQuery {
    prefix: String,
}

impl WithNameStartingWithQuery {
    pub async fn fetch_all(
        &self,
        client: &mut WorkspaceClient,
    ) -> Result<Vec<PgRole>, sqlx::Error> {
        query_as!(
            PgRole,
            r#"
select
    oid as "oid!",
    rolname as "rolname!",
    rolsuper as "rolsuper!",
    rolinherit as "rolinherit!",
    rolcreaterole as "rolcreaterole!",
    rolcreatedb as "rolcreatedb!",
    rolcanlogin as "rolcanlogin!",
    rolreplication as "rolreplication!",
    rolconnlimit as "rolconnlimit!",
    rolvaliduntil,
    rolbypassrls as "rolbypassrls!"
from pg_roles
where starts_with(rolname, $1)
"#,
            self.prefix
        )
        .fetch_all(client.get_conn())
        .await
    }
}

/// A representation of role grants, starting from a single role and traversing
/// the full set (as visible via the current DB connection) of either more
/// specific roles which are granted to it, or more general roles to which it
/// is granted. This is useful, for example, for enumerating permissions which
/// are inherited rather than granted directly to a specific user.
#[derive(Clone, Debug)]
pub struct RoleTree {
    pub role: PgRole,
    pub branches: Vec<RoleTree>,
    pub inherit: bool,
    pub admin: bool,
}

#[derive(Clone, Debug, FromRow)]
struct RoleTreeRow {
    #[sqlx(flatten)]
    role: PgRole,
    parent: Option<Oid>,
    inherit: bool,
    admin: bool,
}

impl RoleTree {
    pub fn members_of_oid(role_oid: Oid) -> MembersOfOidQuery {
        MembersOfOidQuery { role: role_oid }
    }

    pub fn members_of_rolname(rolname: &str) -> MembersOfRolnameQuery {
        MembersOfRolnameQuery {
            role: rolname.to_owned(),
        }
    }

    pub fn granted_to(role_oid: Oid) -> GrantedToQuery {
        GrantedToQuery { role_oid }
    }

    pub fn granted_to_rolname<'a>(rolname: &'a str) -> GrantedToRolnameQuery<'a> {
        GrantedToRolnameQuery { rolname }
    }

    /// Flattens tree structure, omitting any nodes which are not implicitly
    /// inherited by the root role. Each of the resulting item's `branches`
    /// fields is emptied.
    pub fn flatten_inherited(self) -> Vec<Self> {
        [
            vec![Self {
                branches: vec![],
                ..self
            }],
            self.branches
                .into_iter()
                .filter(|member| member.inherit)
                .map(|member| member.flatten_inherited())
                .collect::<Vec<_>>()
                .concat(),
        ]
        .concat()
    }
}

/// Constructs a SQL string to query a tree of Postgres role grants.
///
/// `param_cast` is the SQL required to cast the Postgres input paramter (`$1`)
/// to type OID, for example: `"::regrole::oid"` if `$1` is a rolname or `""` if
/// `$1` is already an OID.
///
/// The query consists of 3 stages:
/// - First, the recursive CTE `memberships_all` constructs the root node and
///   then recurses across Postgres role grants directly or indirectly connected
///   to it. This produces a temporary table with columns:
///   - `roleid`: OID of the role.
///   - `parent`: Reference to the role grant traversed to get to the current
///     node. (OID of the role this node is a member of or is granted to,
///     depending on the traversal direction.)
///   - `admin`: Whether the role grant includes `WITH ADMIN OPTION`. This is
///     set to `false` for the root node, because in Postgres roles are not
///     permitted to grant themselves to others.
///   - `inherit`: Whether there is a contiguous path of `WITH INHERIT OPTION`
///     between this role grant and the starting node (inclusive of both ends).
///     This is set to `true` for the root node, because inheritance of
///     permissions from a role to itself is self-evident.
/// - Next, `memberships_agg` aggregates over `memberships_all` to consolidate
///   role grants which have been performed multiple times by distinct grantors.
///   If *any* of the paths from the root node to a role grant confer the
///   `admin` and/or `inherit` characteristics, then the corresponding property
///   is set to `true`.
/// - Finally, `memberships_agg` is joined to the `pg_roles` view for
///   convenience, to return detailed information on each row beyond its OID.
fn construct_role_tree_query(param_cast: &str, direction: RoleTreeDirection) -> String {
    format!(
        "
with recursive memberships_all as (
    select
        $1{param_cast} as roleid,
        null::oid      as parent,
        false          as admin,
        true           as inherit
    union all
        select
            {roleid_col}                                  as roleid,
            prev.roleid                                   as parent,
            membership.admin_option                       as admin,
            prev.inherit and membership.inherit_option    as inherit
        from memberships_all as prev join pg_auth_members as membership
            on {cte_join_cond}
)
select
    pg_roles.*,
    memberships_agg.parent  as parent,
    memberships_agg.admin   as admin,
    memberships_agg.inherit as inherit
from (
    select
        roleid,
        parent,
        bool_or(admin)   as admin,
        bool_or(inherit) as inherit
    from memberships_all
    group by roleid, parent
) as memberships_agg join pg_roles on pg_roles.oid = memberships_agg.roleid
",
        cte_join_cond = match direction {
            RoleTreeDirection::MembersOf => "membership.roleid = prev.roleid",
            RoleTreeDirection::GrantedTo => "membership.member = prev.roleid",
        },
        roleid_col = match direction {
            RoleTreeDirection::MembersOf => "membership.member",
            RoleTreeDirection::GrantedTo => "membership.roleid",
        },
    )
}

enum RoleTreeDirection {
    MembersOf,
    GrantedTo,
}

#[derive(Clone, Debug)]
pub struct MembersOfOidQuery {
    role: Oid,
}
impl MembersOfOidQuery {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        let rows: Vec<RoleTreeRow> =
            query_as(&construct_role_tree_query("", RoleTreeDirection::MembersOf))
                .bind(self.role)
                .fetch_all(client.get_conn())
                .await?;
        Ok(rows
            .iter()
            .find(|row| row.parent.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                admin: root_row.admin,
                inherit: root_row.inherit,
            }))
    }
}

#[derive(Clone, Debug)]
pub struct MembersOfRolnameQuery {
    role: String,
}

impl MembersOfRolnameQuery {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        let rows: Vec<RoleTreeRow> = query_as(&construct_role_tree_query(
            "::regrole::oid",
            RoleTreeDirection::MembersOf,
        ))
        .bind(self.role.as_str() as &str)
        .fetch_all(client.get_conn())
        .await?;
        Ok(rows
            .iter()
            .find(|row| row.parent.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                admin: root_row.admin,
                inherit: root_row.inherit,
            }))
    }
}

#[derive(Clone, Debug)]
pub struct GrantedToQuery {
    role_oid: Oid,
}

impl GrantedToQuery {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        let rows: Vec<RoleTreeRow> =
            query_as(&construct_role_tree_query("", RoleTreeDirection::GrantedTo))
                .bind(self.role_oid)
                .fetch_all(client.get_conn())
                .await?;
        Ok(rows
            .iter()
            .find(|row| row.parent.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                admin: root_row.admin,
                inherit: root_row.inherit,
            }))
    }
}

#[derive(Clone, Copy, Debug)]
pub struct GrantedToRolnameQuery<'a> {
    rolname: &'a str,
}

impl<'a> GrantedToRolnameQuery<'a> {
    pub async fn fetch_tree(
        self,
        client: &mut WorkspaceClient,
    ) -> Result<Option<RoleTree>, sqlx::Error> {
        let rows: Vec<RoleTreeRow> = query_as(&construct_role_tree_query(
            "::regrole::oid",
            RoleTreeDirection::GrantedTo,
        ))
        .bind(self.rolname)
        .fetch_all(client.get_conn())
        .await?;
        Ok(rows
            .iter()
            .find(|row| row.parent.is_none())
            .map(|root_row| RoleTree {
                role: root_row.role.clone(),
                branches: compute_members(&rows, root_row.role.oid),
                admin: root_row.admin,
                inherit: root_row.inherit,
            }))
    }
}

fn compute_members(rows: &Vec<RoleTreeRow>, root: Oid) -> Vec<RoleTree> {
    rows.iter()
        .filter(|row| row.parent == Some(root))
        .map(|row| RoleTree {
            role: row.role.clone(),
            branches: compute_members(rows, row.role.oid),
            admin: row.admin,
            inherit: row.inherit,
        })
        .collect()
}

#[derive(Debug, Error)]
pub enum RolnameParseError {
    #[error("rolname does not have phono user prefix")]
    MissingPrefix,
    #[error("unable to parse uuid from rolname: {0}")]
    BadUuid(uuid::Error),
}

pub fn user_id_from_rolname(rolname: &str) -> Result<Uuid, RolnameParseError> {
    if rolname.starts_with(ROLE_PREFIX_USER) {
        let mut rolname = rolname.to_owned();
        rolname.replace_range(0..ROLE_PREFIX_USER.len(), "");
        Uuid::parse_str(&rolname).map_err(RolnameParseError::BadUuid)
    } else {
        Err(RolnameParseError::MissingPrefix)
    }
}
