import os
from smolagents import Tool
from huggingface_hub import list_datasets,list_models,list_organization_members,list_collections,list_spaces

class AtlasiaHubInfoTool(Tool):
    name = "atlasia_huggingface_hub_search"
    description = "Search and retrieve information about datasets, models, and spaces of atlasia organization from the HuggingFace Hub"
    inputs = {
        "resource_type": {
            "type": "string",
            "description": "Type of resource to search for. Must be one of: 'dataset', 'model', 'space', 'collection', or 'organization_members'",
            "enum": ["dataset", "model", "space", "collection", "organization_members"]
        }
    }
       
    output_type="array"

    def forward(self,resource_type:str)-> list:
        if not isinstance(resource_type,str):
            raise ValueError("resource_type should be a string")
        if resource_type=="dataset":
            return list(list_datasets(author="atlasia"))
        elif resource_type=="model":
            return list(list_models(author="atlasia"))
        elif resource_type=="space":
            return list(list_spaces(author="atlasia"))
        elif resource_type=="collection":
            return list(list_collections(owner="atlasia"))
        elif resource_type=="organization_members":
            return [x.fullname for x in list(list_organization_members(organization="atlasia"))]
        else:
            raise ValueError("resource_type should be one of: 'dataset', 'model', 'space', 'collection', or 'organization_members'")