use itertools::Itertools;
use unitycatalog_common::models::PropertyMap;
use unitycatalog_common::store::{ResourceStore, ResourceStoreReader};
use unitycatalog_common::{
    AssociationLabel, EMPTY_RESOURCE_NAME, Object, ObjectLabel, Resource, ResourceIdent,
    ResourceName, ResourceRef,
};
use unitycatalog_common::{Error, Result};

use crate::GraphStore;

pub trait IdentRefs {
    fn ident(&self) -> (&ObjectLabel, &ResourceRef);
}

impl IdentRefs for ResourceIdent {
    fn ident(&self) -> (&ObjectLabel, &ResourceRef) {
        (self.as_ref(), self.as_ref())
    }
}

#[async_trait::async_trait]
impl ResourceStoreReader for GraphStore {
    /// Get a resource by its identifier.
    ///
    /// # Arguments
    /// - `id`: The identifier of the resource to get.
    ///
    /// # Returns
    /// The resource with the given identifier.
    async fn get(&self, id: &ResourceIdent) -> Result<(Resource, ResourceRef)> {
        match id.as_ref() {
            ResourceRef::Uuid(uuid) => Ok((self.get_object(uuid).await?.try_into()?, id.into())),
            ResourceRef::Name(name) => {
                let object = self.get_object_by_name(id.as_ref(), name).await?;
                let id_new = ResourceRef::Uuid(object.id);
                Ok((object.try_into()?, id_new))
            }
            ResourceRef::Undefined => Err(Error::generic("Cannot get undefined resource")),
        }
    }

    /// List resources.
    ///
    /// List resources in the store that are children of the given resource.
    /// If the Reference inside the ResourceIdent is [Undefined](crate::ResourceRef::Undefined),
    /// the root of the store is used and resources of the specified type are listed.
    ///
    /// # Arguments
    /// - `root`: The root resource to list children of.
    /// - `max_results`: The maximum number of results to return.
    /// - `page_token`: The token to use to get the next page of results.
    async fn list(
        &self,
        label: &ObjectLabel,
        namespace: Option<&ResourceName>,
        max_results: Option<usize>,
        page_token: Option<String>,
    ) -> Result<(Vec<Resource>, Option<String>)> {
        let namespace = namespace.unwrap_or_else(|| &EMPTY_RESOURCE_NAME);
        let objects = self
            .list_objects(label, namespace, page_token.as_deref(), max_results)
            .await?;
        Ok((
            objects
                .0
                .into_iter()
                .map(|object| object.try_into())
                .try_collect()?,
            objects.1,
        ))
    }
}

#[async_trait::async_trait]
impl ResourceStore for GraphStore {
    /// Create a new resource.
    ///
    /// # Arguments
    /// - `resource`: The resource to create.
    ///
    /// # Returns
    /// The created resource.
    async fn create(&self, resource: Resource) -> Result<(Resource, ResourceRef)> {
        let object: Object = resource.try_into()?;
        // A non-nil id means the caller pre-allocated it (e.g. a managed volume or
        // staging table that embeds the id in its storage path); a nil id lets the
        // store mint a fresh v7. API request types carry no id field, so callers
        // cannot force an id through this path.
        let supplied_id = (!object.id.is_nil()).then_some(object.id);
        let object = self
            .add_object(&object.label, &object.name, object.properties, supplied_id)
            .await?;
        let id = ResourceRef::Uuid(object.id);
        Ok((object.try_into()?, id))
    }

    /// Delete a resource by its identifier.
    ///
    /// # Arguments
    /// - `id`: The identifier of the resource to delete.
    async fn delete(&self, id: &ResourceIdent) -> Result<()> {
        match id.as_ref() {
            ResourceRef::Uuid(uuid) => self.delete_object(uuid).await?,
            ResourceRef::Name(name) => {
                let obj = self.get_object_by_name(id.as_ref(), name).await?;
                self.delete_object(&obj.id).await?;
            }
            ResourceRef::Undefined => {
                return Err(Error::generic("Cannot delete undefined resource"));
            }
        };
        Ok(())
    }

    /// Update a resource.
    ///
    /// # Arguments
    /// - `resource`: The resource to update.
    ///
    /// # Returns
    /// The updated resource.
    async fn update(
        &self,
        id: &ResourceIdent,
        resource: Resource,
    ) -> Result<(Resource, ResourceRef)> {
        let (uuid, _object) = self.ident_to_uuid(id).await?;
        let object: Object = resource.try_into()?;
        Ok((
            self.update_object(
                &uuid,
                Some(&object.label),
                Some(object.name.as_ref()),
                object.properties,
            )
            .await?
            .try_into()?,
            uuid.into(),
        ))
    }

    /// Add an association between two resources.
    async fn add_association(
        &self,
        from: &ResourceIdent,
        to: &ResourceIdent,
        label: &AssociationLabel,
        properties: Option<PropertyMap>,
    ) -> Result<()> {
        let from_id = match from.as_ref() {
            ResourceRef::Uuid(uuid) => *uuid,
            ResourceRef::Name(name) => self.get_object_by_name(from.as_ref(), name).await?.id,
            ResourceRef::Undefined => {
                return Err(Error::generic(
                    "Cannot add association to undefined resource",
                ));
            }
        };
        let to_id = match to.as_ref() {
            ResourceRef::Uuid(uuid) => *uuid,
            ResourceRef::Name(name) => self.get_object_by_name(to.as_ref(), name).await?.id,
            ResourceRef::Undefined => {
                return Err(Error::generic(
                    "Cannot add association to undefined resource",
                ));
            }
        };
        self.add_association(
            &from_id,
            label,
            &to_id,
            properties.map(|p| serde_json::Value::Object(p.into_iter().collect())),
        )
        .await?;
        Ok(())
    }

    async fn remove_association(
        &self,
        _from: &ResourceIdent,
        _to: &ResourceIdent,
        _label: &AssociationLabel,
    ) -> Result<()> {
        todo!("remove_association")
    }

    async fn list_associations(
        &self,
        resource: &ResourceIdent,
        label: &AssociationLabel,
        target_label: Option<&ResourceIdent>,
        max_results: Option<usize>,
        page_token: Option<String>,
    ) -> Result<(Vec<ResourceIdent>, Option<String>)> {
        let target_label = target_label.map(|r| r.ident().0);
        let (target_id, _) = self.ident_to_uuid(resource).await?;
        let (associations, token) = self
            .list_associations(
                &target_id,
                label,
                target_label,
                page_token.as_deref(),
                max_results,
            )
            .await?;
        let idents = associations
            .into_iter()
            .map(|assoc| assoc.to_label.to_ident(assoc.to_id))
            .collect();
        Ok((idents, token))
    }

    async fn list_associations_with_properties(
        &self,
        resource: &ResourceIdent,
        label: &AssociationLabel,
        target_label: Option<&ResourceIdent>,
        max_results: Option<usize>,
        page_token: Option<String>,
    ) -> Result<(Vec<(ResourceIdent, Option<PropertyMap>)>, Option<String>)> {
        let target_label = target_label.map(|r| r.ident().0);
        let (target_id, _) = self.ident_to_uuid(resource).await?;
        let (associations, token) = self
            .list_associations(
                &target_id,
                label,
                target_label,
                page_token.as_deref(),
                max_results,
            )
            .await?;
        let entries = associations
            .into_iter()
            .map(|assoc| {
                let props = match assoc.properties {
                    Some(serde_json::Value::Object(map)) => Some(map.into_iter().collect()),
                    _ => None,
                };
                (assoc.to_label.to_ident(assoc.to_id), props)
            })
            .collect();
        Ok((entries, token))
    }
}
