//! Storage layer for managing objects and associations.
//!
//! Lossely based on the data model applied in Meta's [TAO].
//!
//! [TAO]: https://www.usenix.org/system/files/conference/atc13/atc13-bronson.pdf

use sqlx::PgPool;
use sqlx::migrate::Migrator;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
use unitycatalog_common::services::encryption::EnvelopeEncryptor;
use unitycatalog_common::{ResourceIdent, ResourceRef};
use uuid::Uuid;

use super::{Association, AssociationLabel, Object, ObjectLabel};
use crate::constants::MAX_PAGE_SIZE;
use crate::pagination::V1PaginateToken;
use crate::resources::IdentRefs as _;
use crate::{error::Result, pagination::PaginateToken};

static MIGRATOR: Migrator = sqlx::migrate!();

#[derive(Clone)]
pub struct Store {
    pub(crate) pool: PgPool,
    pub(crate) encryptor: EnvelopeEncryptor,
}

impl Store {
    pub fn new(pool: PgPool, encryptor: EnvelopeEncryptor) -> Self {
        Self { pool, encryptor }
    }

    pub async fn connect(url: impl AsRef<str>, encryptor: EnvelopeEncryptor) -> Result<Self> {
        let options: PgConnectOptions = url.as_ref().parse()?;
        let pool_options = PgPoolOptions::new().max_connections(96);
        let pool = pool_options.connect_with(options).await?;
        Ok(Self::new(pool, encryptor))
    }

    pub async fn migrate(&self) -> Result<()> {
        MIGRATOR.run(&self.pool).await?;
        Ok(())
    }

    /// Convert a resource reference to a UUID.
    ///
    /// If the reference is a name, the corresponding object is fetched from the store.
    /// to get the UUID. The object is returned as well in case it is needed later
    /// to avoid an additional fetch.
    ///
    /// # Parameters
    /// - `reference`: The reference to convert.
    ///
    /// # Returns
    /// The UUID of the reference and the object if the reference is a name.
    ///
    /// # Errors
    /// In case of an undefined reference, an error is returned.
    pub async fn ident_to_uuid(&self, reference: &ResourceIdent) -> Result<(Uuid, Option<Object>)> {
        let (label, ident) = reference.ident();
        match ident {
            ResourceRef::Uuid(id) => Ok((*id, None)),
            ResourceRef::Name(name) => {
                let object = self.get_object_by_name(label, name).await?;
                Ok((object.id, Some(object)))
            }
            ResourceRef::Undefined => Err(crate::Error::entity_not_found("undefined")),
        }
    }

    /// Add an object to the store.
    ///
    /// # Parameters
    /// - `label`: The label of the object.
    /// - `name`: The namespaced name of the object.
    /// - `name`: The name of the object.
    /// - `properties`: The properties of the object.
    ///
    /// # Returns
    /// The object that was added to the store.
    ///
    /// # Errors
    /// - [AlreadyExists](crate::Error::AlreadyExists): If an object with the
    ///   same name already exists in the namespace
    pub async fn add_object(
        &self,
        label: &ObjectLabel,
        name: &[String],
        properties: Option<serde_json::Value>,
        id: Option<Uuid>,
    ) -> Result<Object> {
        let mut txn = self.pool.begin().await?;
        let obj = add_object(label, name, properties, id, &mut txn).await?;
        txn.commit().await?;
        Ok(obj)
    }

    /// Get an object from the store.
    ///
    /// ## Parameters
    /// - `id`: The globally unique identifier of the object.
    ///
    /// ## Returns
    /// The object with the given identifier.
    ///
    /// ## Errors
    /// - [EntityNotFound](crate::Error::EntityNotFound): If the object does not exist.
    pub async fn get_object(&self, id: &Uuid) -> Result<Object> {
        let mut conn = self.pool.acquire().await?;
        let obj = get_object(id, &mut conn).await?;
        Ok(obj)
    }

    /// Get an object from the store by name.
    ///
    /// The name of the object is unique within the namespace.
    ///
    /// ## Parameters
    /// - `label`: The label of the object.
    /// - `namespace`: The namespace of the object.
    /// - `name`: The name of the object.
    ///
    /// ## Returns
    /// The object with the given name.
    ///
    /// ## Errors
    /// - [EntityNotFound](crate::Error::EntityNotFound): If the object does not exist.
    pub async fn get_object_by_name(&self, label: &ObjectLabel, name: &[String]) -> Result<Object> {
        let mut conn = self.pool.acquire().await?;
        let obj = get_object_by_name(label, name, &mut conn).await?;
        Ok(obj)
    }

    /// Update an object in the store.
    ///
    /// ## Parameters
    /// - `id`: The globally unique identifier of the object.
    /// - `properties`: The properties of the object.
    ///
    /// ## Returns
    /// The updated object.
    ///
    /// ## Errors
    /// - [EntityNotFound](crate::Error::EntityNotFound): If the object does not exist.
    pub async fn update_object(
        &self,
        id: &Uuid,
        new_label: impl Into<Option<&ObjectLabel>>,
        new_name: impl Into<Option<&[String]>>,
        properties: impl Into<Option<serde_json::Value>>,
    ) -> Result<Object> {
        let mut txn = self.pool.begin().await?;
        let obj = update_object(id, new_label, new_name, properties, &mut txn).await?;
        txn.commit().await?;
        Ok(obj)
    }

    /// Delete an object from the store.
    ///
    /// ## Parameters
    /// - `id`: The globally unique identifier of the object.
    pub async fn delete_object(&self, id: &Uuid) -> Result<()> {
        let mut txn = self.pool.begin().await?;
        delete_object(id, &mut txn).await?;
        Ok(txn.commit().await?)
    }

    /// List objects from the store.
    ///
    /// Returns a list of objects in the namespace. The list is paginated.
    ///
    /// ## Parameters
    /// - `label`: The label of the objects.
    /// - `namespace`: The namespace of the objects.
    /// - `page_token`: The page token.
    /// - `max_page_size`: The maximum page size.
    ///
    /// ## Returns
    /// A tuple containing the objects in the namespace and an optional next page token.
    pub async fn list_objects(
        &self,
        label: &ObjectLabel,
        namespace: &[String],
        page_token: Option<&str>,
        max_page_size: Option<usize>,
    ) -> Result<(Vec<Object>, Option<String>)> {
        let max_page_size = usize::min(max_page_size.unwrap_or(MAX_PAGE_SIZE), MAX_PAGE_SIZE);
        let token = page_token
            .map(PaginateToken::<Uuid>::try_from)
            .transpose()?;
        let (_token_ts, token_id) = token
            .as_ref()
            .map(|PaginateToken::V1(V1PaginateToken { created_at, id })| (created_at, id))
            .unzip();
        let mut conn = self.pool.acquire().await?;
        list_objects(label, namespace, max_page_size, token_id, &mut conn).await
    }

    /// Add an association to the store.
    ///
    /// Associations are directed edges between objects.
    /// If an inverse association exists, it is automatically created.
    ///
    /// ## Parameters
    /// - `from_id`: The identifier of the source object.
    /// - `label`: The label of the association.
    /// - `to_id`: The identifier of the target object.
    /// - `properties`: The properties of the association.
    ///
    /// ## Returns
    /// The association that was added to the store.
    ///
    /// ## Errors
    /// - [EntityNotFound](crate::Error::EntityNotFound): If the source or target object does not exist.
    /// - [AlreadyExists](crate::Error::AlreadyExists): If the association already exists.
    pub async fn add_association(
        &self,
        from_id: &Uuid,
        label: &AssociationLabel,
        to_id: &Uuid,
        properties: impl Into<Option<serde_json::Value>>,
    ) -> Result<Association> {
        let mut txn = self.pool.begin().await?;
        let association = add_association(from_id, label, to_id, properties, &mut txn).await?;
        txn.commit().await?;
        Ok(association)
    }

    /// Delete an association from the store.
    ///
    /// If an inverse association exists, it is automatically deleted.
    pub async fn delete_association(
        &self,
        from_id: &Uuid,
        label: &AssociationLabel,
        to_id: &Uuid,
    ) -> Result<()> {
        let mut txn = self.pool.begin().await?;
        delete_association(from_id, label, to_id, &mut txn).await?;
        txn.commit().await?;
        Ok(())
    }

    /// List associations of a specific type from an object to a set of objects.
    ///
    /// ## Parameters
    /// - `from_id`: The identifier of the source object.
    /// - `label`: The label of the association.
    /// - `to_ids`: The identifiers of the target objects.
    ///
    /// ## Returns
    /// The associations from the source object to the target objects.
    pub async fn get_associations(
        &self,
        from_id: &Uuid,
        label: &AssociationLabel,
        to_ids: &[Uuid],
        page_token: Option<&str>,
        max_page_size: Option<usize>,
    ) -> Result<(Vec<Association>, Option<String>)> {
        let max_page_size = usize::min(max_page_size.unwrap_or(MAX_PAGE_SIZE), MAX_PAGE_SIZE);
        let token = page_token
            .map(PaginateToken::<Uuid>::try_from)
            .transpose()?;
        let (_token_ts, token_id) = token
            .as_ref()
            .map(|PaginateToken::V1(V1PaginateToken { created_at, id })| (created_at, id))
            .unzip();
        let mut conn = self.pool.acquire().await?;
        get_associations(from_id, label, to_ids, max_page_size, token_id, &mut conn).await
    }

    /// List associations of a specific type from an object to all objects.
    pub async fn list_associations(
        &self,
        from_id: &Uuid,
        label: &AssociationLabel,
        target_label: Option<&ObjectLabel>,
        page_token: Option<&str>,
        max_page_size: Option<usize>,
    ) -> Result<(Vec<Association>, Option<String>)> {
        let max_page_size = usize::min(max_page_size.unwrap_or(MAX_PAGE_SIZE), MAX_PAGE_SIZE);
        let token = page_token
            .map(PaginateToken::<Uuid>::try_from)
            .transpose()?;
        let (_token_ts, token_id) = token
            .as_ref()
            .map(
                |PaginateToken::V1(V1PaginateToken { created_at, id }): &PaginateToken<Uuid>| {
                    (created_at, id)
                },
            )
            .unzip();
        let mut conn = self.pool.acquire().await?;
        list_associations(
            from_id,
            label,
            target_label,
            max_page_size,
            token_id,
            &mut conn,
        )
        .await
    }
}

// --- ObjectStore<ObjectLabel> implementation ---

use olai_store::EMPTY_RESOURCE_NAME;
use olai_store::name::ResourceName;

#[async_trait::async_trait]
impl olai_store::ObjectStoreReader<ObjectLabel> for Store {
    async fn get(&self, id: &Uuid) -> olai_store::Result<olai_store::Object<ObjectLabel>> {
        Ok(self.get_object(id).await?)
    }

    async fn get_by_name(
        &self,
        label: ObjectLabel,
        name: &ResourceName,
    ) -> olai_store::Result<olai_store::Object<ObjectLabel>> {
        Ok(self.get_object_by_name(&label, name).await?)
    }

    async fn list(
        &self,
        label: ObjectLabel,
        namespace: Option<&ResourceName>,
        max_results: Option<usize>,
        page_token: Option<String>,
    ) -> olai_store::Result<(Vec<olai_store::Object<ObjectLabel>>, Option<String>)> {
        let namespace = namespace.unwrap_or(&EMPTY_RESOURCE_NAME);
        let (objects, token) = self
            .list_objects(&label, namespace, page_token.as_deref(), max_results)
            .await?;
        Ok((objects, token))
    }
}

#[async_trait::async_trait]
impl olai_store::ObjectStore<ObjectLabel> for Store {
    async fn create(
        &self,
        label: ObjectLabel,
        name: &ResourceName,
        properties: Option<serde_json::Value>,
        id: Option<Uuid>,
    ) -> olai_store::Result<olai_store::Object<ObjectLabel>> {
        Ok(self.add_object(&label, name, properties, id).await?)
    }

    async fn update(
        &self,
        id: &Uuid,
        properties: Option<serde_json::Value>,
    ) -> olai_store::Result<olai_store::Object<ObjectLabel>> {
        Ok(self
            .update_object(id, None::<&ObjectLabel>, None::<&[String]>, properties)
            .await?)
    }

    async fn delete(&self, id: &Uuid) -> olai_store::Result<()> {
        Ok(self.delete_object(id).await?)
    }
}

// --- AssociationStore<ObjectLabel> implementation ---

#[async_trait::async_trait]
impl olai_store::AssociationStoreReader<ObjectLabel> for Store {
    async fn list(
        &self,
        from_id: Uuid,
        label: &str,
        target_label: Option<ObjectLabel>,
        max_results: Option<usize>,
        page_token: Option<String>,
    ) -> olai_store::Result<(Vec<olai_store::Association<ObjectLabel>>, Option<String>)> {
        let assoc_label: AssociationLabel = label.parse().map_err(|_| {
            olai_store::Error::InvalidArgument(format!("Unknown association label: {label}"))
        })?;
        let (associations, token) = self
            .list_associations(
                &from_id,
                &assoc_label,
                target_label.as_ref(),
                page_token.as_deref(),
                max_results,
            )
            .await?;
        let converted: Vec<olai_store::Association<ObjectLabel>> = associations
            .into_iter()
            .map(|a| olai_store::Association {
                id: a.id,
                from_id: a.from_id,
                label: a.label.to_string(),
                to_id: a.to_id,
                to_label: a.to_label,
                properties: a.properties,
                created_at: a.created_at,
                updated_at: a.updated_at,
            })
            .collect();
        Ok((converted, token))
    }
}

#[async_trait::async_trait]
impl olai_store::AssociationStore<ObjectLabel> for Store {
    async fn add(
        &self,
        from_id: Uuid,
        to_id: Uuid,
        label: &str,
        properties: Option<serde_json::Value>,
    ) -> olai_store::Result<()> {
        let assoc_label: AssociationLabel = label.parse().map_err(|_| {
            olai_store::Error::InvalidArgument(format!("Unknown association label: {label}"))
        })?;
        self.add_association(&from_id, &assoc_label, &to_id, properties)
            .await?;
        Ok(())
    }

    async fn remove(&self, from_id: Uuid, to_id: Uuid, label: &str) -> olai_store::Result<()> {
        let assoc_label: AssociationLabel = label.parse().map_err(|_| {
            olai_store::Error::InvalidArgument(format!("Unknown association label: {label}"))
        })?;
        self.delete_association(&from_id, &assoc_label, &to_id)
            .await?;
        Ok(())
    }
}

async fn list_associations(
    from_id: &Uuid,
    label: &AssociationLabel,
    target_label: Option<&ObjectLabel>,
    max_page_size: usize,
    token_id: Option<&Uuid>,
    conn: &mut sqlx::pool::PoolConnection<sqlx::Postgres>,
) -> std::result::Result<(Vec<Association>, Option<String>), crate::Error> {
    let assocs = sqlx::query_as!(
        Association,
        r#"
                SELECT
                    id,
                    from_id,
                    label AS "label: AssociationLabel",
                    to_id,
                    properties,
                    created_at,
                    updated_at,
                    to_label as "to_label: ObjectLabel"
                FROM associations
                WHERE from_id = $1
                  AND label = $2
                  AND ( to_label = $3 OR $3 IS NULL )
                  -- Pagination
                  AND ( id < $4 OR $4 IS NULL )
                ORDER BY id DESC
                LIMIT $5
                "#,
        from_id,
        label as &AssociationLabel,
        target_label as Option<&ObjectLabel>,
        token_id,
        max_page_size as i64
    )
    .fetch_all(&mut **conn)
    .await?;

    let next = (assocs.len() == max_page_size)
        .then(|| {
            assocs.last().map(|a| {
                PaginateToken::V1(V1PaginateToken {
                    created_at: a.created_at,
                    id: a.id,
                })
                .to_string()
            })
        })
        .flatten();

    Ok((assocs, next))
}

async fn add_object(
    label: &ObjectLabel,
    name: &[String],
    properties: Option<serde_json::Value>,
    id: Option<Uuid>,
    txn: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<Object, crate::Error> {
    // A caller that pre-allocates an id (e.g. a managed volume or staging table
    // embedding the id in its storage path) passes it here; `COALESCE` falls back
    // to the column's `uuidv7()` default for the common id-less create.
    let obj = sqlx::query_as!(
        Object,
        r#"
                INSERT INTO objects ( id, label, name, properties )
                VALUES ( COALESCE($1, uuidv7()), $2, $3, $4 )
                RETURNING
                    id,
                    label AS "label: ObjectLabel",
                    name,
                    properties,
                    created_at,
                    updated_at
                "#,
        id,
        label as &ObjectLabel,
        name,
        properties
    )
    .fetch_one(&mut **txn)
    .await?;
    Ok(obj)
}

async fn get_object(
    id: &Uuid,
    conn: &mut sqlx::pool::PoolConnection<sqlx::Postgres>,
) -> Result<Object, crate::Error> {
    let obj = sqlx::query_as!(
        Object,
        r#"
                SELECT
                    id,
                    label AS "label: ObjectLabel",
                    name,
                    properties,
                    created_at,
                    updated_at
                FROM objects
                WHERE id = $1
                "#,
        id
    )
    .fetch_one(&mut **conn)
    .await?;
    Ok(obj)
}

async fn get_object_by_name(
    label: &ObjectLabel,
    name: &[String],
    conn: &mut sqlx::pool::PoolConnection<sqlx::Postgres>,
) -> Result<Object, crate::Error> {
    let obj = sqlx::query_as!(
        Object,
        r#"
                SELECT
                    id,
                    label AS "label: ObjectLabel",
                    name,
                    properties,
                    created_at,
                    updated_at
                FROM objects
                WHERE label = $1
                  AND name = $2
                "#,
        label as &ObjectLabel,
        name
    )
    .fetch_one(&mut **conn)
    .await?;
    Ok(obj)
}

async fn update_object(
    id: &Uuid,
    new_label: impl Into<Option<&ObjectLabel>>,
    new_name: impl Into<Option<&[String]>>,
    properties: impl Into<Option<serde_json::Value>>,
    txn: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<Object, crate::Error> {
    let obj = sqlx::query_as!(
        Object,
        r#"
                UPDATE objects
                SET
                    label = COALESCE($2, label),
                    name = COALESCE($3, name),
                    properties = COALESCE($4, properties)
                WHERE id = $1
                RETURNING
                    id,
                    label AS "label: ObjectLabel",
                    name,
                    properties,
                    created_at,
                    updated_at
                "#,
        id,
        new_label.into() as Option<&ObjectLabel>,
        new_name.into(),
        properties.into()
    )
    .fetch_one(&mut **txn)
    .await?;
    Ok(obj)
}

async fn delete_object(
    id: &Uuid,
    txn: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<(), crate::Error> {
    sqlx::query!(
        r#"
                DELETE FROM associations
                WHERE from_id = $1 OR to_id = $1
                "#,
        id
    )
    .execute(&mut **txn)
    .await?;
    sqlx::query!(
        r#"
                DELETE FROM objects
                WHERE id = $1
                "#,
        id
    )
    .execute(&mut **txn)
    .await?;
    Ok(())
}

async fn list_objects(
    label: &ObjectLabel,
    namespace: &[String],
    max_page_size: usize,
    token_id: Option<&Uuid>,
    conn: &mut sqlx::pool::PoolConnection<sqlx::Postgres>,
) -> std::result::Result<(Vec<Object>, Option<String>), crate::Error> {
    let objects = sqlx::query_as!(
        Object,
        r#"
                SELECT
                    id,
                    label AS "label: ObjectLabel",
                    name,
                    properties,
                    created_at,
                    updated_at
                FROM objects
                WHERE label = $1
                    AND ( $2 = 0 OR name[1:$2] = $3)
                    AND ( id < $4 OR $4 IS NULL )
                ORDER BY id DESC
                LIMIT $5
                "#,
        label as &ObjectLabel,
        namespace.len() as i32,
        namespace,
        token_id,
        max_page_size as i64
    )
    .fetch_all(&mut **conn)
    .await?;

    let next = (objects.len() == max_page_size)
        .then(|| {
            objects.last().map(|o| {
                PaginateToken::V1(V1PaginateToken {
                    created_at: o.created_at,
                    id: o.id,
                })
                .to_string()
            })
        })
        .flatten();

    Ok((objects, next))
}

async fn get_associations(
    from_id: &Uuid,
    label: &AssociationLabel,
    to_ids: &[Uuid],
    max_page_size: usize,
    token_id: Option<&Uuid>,
    conn: &mut sqlx::pool::PoolConnection<sqlx::Postgres>,
) -> std::result::Result<(Vec<Association>, Option<String>), crate::Error> {
    let assocs = sqlx::query_as!(
        Association,
        r#"
                SELECT
                    id,
                    from_id,
                    label AS "label: AssociationLabel",
                    to_id,
                    properties,
                    created_at,
                    updated_at,
                    to_label as "to_label: ObjectLabel"
                FROM associations
                WHERE from_id = $1
                  AND label = $2
                  AND to_id = ANY($3)
                  -- Pagination
                  AND ( id < $4 OR $4 IS NULL )
                ORDER BY id DESC
                LIMIT $5
                "#,
        from_id,
        label as &AssociationLabel,
        &to_ids,
        token_id,
        max_page_size as i64
    )
    .fetch_all(&mut **conn)
    .await?;

    let next = (assocs.len() == max_page_size)
        .then(|| {
            assocs.last().map(|a| {
                PaginateToken::V1(V1PaginateToken {
                    created_at: a.created_at,
                    id: a.id,
                })
                .to_string()
            })
        })
        .flatten();

    Ok((assocs, next))
}

pub async fn add_association(
    from_id: &Uuid,
    label: &AssociationLabel,
    to_id: &Uuid,
    properties: impl Into<Option<serde_json::Value>>,
    txn: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<Association> {
    let properties = properties.into();
    let to_label = sqlx::query!(
        r#"
            SELECT id, label AS "label: ObjectLabel"
            FROM objects
            WHERE id = $1 OR id = $2
            "#,
        from_id,
        to_id
    )
    .fetch_all(&mut **txn)
    .await?;

    let id_map = to_label
        .into_iter()
        .map(|o| (o.id, o.label))
        .collect::<std::collections::HashMap<_, _>>();
    let to_label = id_map
        .get(to_id)
        .ok_or(crate::Error::entity_not_found("to_id"))?;
    let from_label = id_map
        .get(from_id)
        .ok_or(crate::Error::entity_not_found("from_id"))?;

    // Add the association.
    let association = sqlx::query_as!(
        Association,
        r#"
            INSERT INTO associations ( from_id, label, to_id, to_label, properties )
            VALUES ( $1, $2, $3, $4, $5 )
            RETURNING
                id,
                from_id,
                label AS "label: AssociationLabel",
                to_id,
                to_label as "to_label: ObjectLabel",
                properties,
                created_at,
                updated_at
            "#,
        from_id,
        label as &AssociationLabel,
        to_id,
        to_label as &ObjectLabel,
        properties.clone()
    )
    .fetch_one(&mut **txn)
    .await?;

    // Add the inverse association.
    if let Some(inverse_label) = label.inverse() {
        sqlx::query!(
            r#"
                INSERT INTO associations ( from_id, label, to_id, to_label, properties )
                VALUES ( $1, $2, $3, $4, $5 )
                "#,
            to_id,
            inverse_label as AssociationLabel,
            from_id,
            from_label as &ObjectLabel,
            properties
        )
        .execute(&mut **txn)
        .await?;
    }

    Ok(association)
}

async fn delete_association(
    from_id: &Uuid,
    label: &AssociationLabel,
    to_id: &Uuid,
    txn: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<(), crate::error::Error> {
    sqlx::query!(
        r#"
            DELETE FROM associations
            WHERE from_id = $1 AND label = $2 AND to_id = $3
            "#,
        from_id,
        label as &AssociationLabel,
        to_id
    )
    .execute(&mut **txn)
    .await?;
    if let Some(inverse_label) = label.inverse() {
        sqlx::query!(
            r#"
                DELETE FROM associations
                WHERE from_id = $1 AND label = $2 AND to_id = $3
                "#,
            to_id,
            inverse_label as AssociationLabel,
            from_id
        )
        .execute(&mut **txn)
        .await?;
    };
    Ok(())
}
