import os
import gradio as gr
from typing import Optional
from huggingface_hub import HfApi, list_models

DUMMY_REFRESH = "Load data from user profile" # Selecting this entry will trigger refresh all.
isdummy = lambda name, dummy = DUMMY_REFRESH: name == dummy or dummy in name

WILD_REFRESH = "**"
iswild = lambda name, wild = WILD_REFRESH: name == wild or wild in name

MYSPACE = os.getenv("SPACE_ID", "")
MYUSER = MYSPACE.split("/")[0]

api = HfApi()

def get_user_repos(username: str, token: str) -> list[str]:
    """Fetch all repos under a user's namespace."""
    # api = HfApi(token=token)
    # repos = api.list_repos_objs(author=username)
    repos = list_models(author=username)
    # repos = [repo.id.split("/", 1)[1] for repo in repos] # Actually need authorid for exists to work.
    repos = sorted([repo.id for repo in repos]) # author/repo.
    return repos

def get_gated_repos(username: str, token: str) -> list[str]:
    """Fetch all gated repos."""
    repos = list_models(author=username, gated = True)
    # repos = [repo.id.split("/", 1)[1] for repo in repos] # Actually need authorid for exists to work.
    repos = sorted([repo.id for repo in repos]) # author/repo.
    return repos

def get_organisation(username: str) -> list[str]:
    """Returns all members in organisation.
    If not an organisation, returns [user].
    """
    try:
        return [u.username for u in api.list_organization_members(username)]
    except Exception as e:
        if isinstance(username, list):
            return username
        else:
            return [username]

def grant_access_to_repo(token, repo_id, username):
    """
    Grant access to a user/organization for a specific repo
    """
    try:
        api.grant_access(repo_id=repo_id, user=username, token=token)
        return f"✅ {repo_id} | {username}: Granted."
    except Exception as e:
        if "already has access to the repo" in str(e):
            return f"🔓 {repo_id} | {username}: Already has access."
        return f"❌ {repo_id} | {username}: Error. {str(e)}"

def grant_access_multiple_repos(token, repo_list, username):
    """Grant access to multiple repositories + users at once.
    Users are dereferenced by organisation.
    """
    if not token:
        return "ERROR: Cannot grant access without a write token."
    if not repo_list:
        return "ERROR: No repos selected."
    if not username:
        return "ERROR: No user selected."
        
    if not isinstance(repo_list, list):
        repo_list = [repo_list]
    lusers = get_organisation(username)
    
    results = []
    for user in lusers:
        for repo_id in repo_list:
            results.append(grant_access_to_repo(token, repo_id, user))
    
    return "\n".join(results)

with gr.Blocks(title="HF Bouncer") as demo:
    gr.Markdown("# ⛹️ Huggingface Bouncer")
    gr.Markdown("Grant access to users/organizations for multiple repositories")

    hf_user = gr.Textbox(
        label = "Huggingface user",
        value = MYUSER,
        info = "User / org whose repos are granted.",
        max_lines = 1,
    )
    hf_token = gr.Textbox(
        label="Huggingface token",
        placeholder="hf_xxxxxxxxxxxxxxxx",
        type="password",
        info="Your HF token with write permissions",
        max_lines=1
    )

    refresh_repos_btn = gr.Button( # Press this button to get external user repo data, after login.
        "🔄 Refresh external data - presets and duplicate repos",
        variant="secondary",
        scale=0,
        min_width=40
    )
    
    with gr.Row():
        with gr.Column():
            multi_repo_input = gr.Dropdown(
                label="Repositories to grant",
                multiselect=True,
                allow_custom_value=True, # Can type other people's repos.
                interactive=True,
                choices=[DUMMY_REFRESH],
                value=[],
                info="Username/repo-name"
            )
            multi_user_input = gr.Textbox(
                label="User/Organization", 
                placeholder="username-or-org-name"
            )
            multi_btn = gr.Button("Grant Access to All", variant="primary")
        
        with gr.Column():
            multi_output = gr.Textbox(label="Results", lines=10)

    # ---Actions---

    # Refresh external data.
    def refresh_all_data(username, token, force = False, indwild = False):
        try:
            # Doesn't matter, listing works without tokens.
            # if not token:
            #     raise ValueError("Please login first!")
            # Refresh repos
            lrepos = get_gated_repos(username, token)
            # Add wildcard at the top.
            lreposw = [WILD_REFRESH] + lrepos

            # Check if wildcard selected, select all values (besides wc).
            lsel = []
            if indwild: # Wildcard selected, 
                lsel = lrepos
            return gr.Dropdown(choices=lreposw, value=lsel)
        except Exception as e:
            gr.Warning(f"Failed to refresh data: {str(e)}")
            return gr.Dropdown(choices=[DUMMY_REFRESH], value=[])

    def force_refresh(username, token):
        """Run refresh with force.
        Lambda method doesn't seem to work for some reason,
        perhaps profile and oauth aren't functional with args/kwargs.
        """
        return refresh_all_data(username, token, force = True)
    
    refresh_repos_btn.click(
        fn=force_refresh,
        inputs=[hf_user, hf_token],
        outputs=[multi_repo_input]
    )

    def check_dummy(username, token, dropdown_value = None):
        """If dummy was selected, initiate dropdown refresh.
        
        """
        name = ""
        if not dropdown_value: # Empty list or null.
            pass
        elif isinstance(dropdown_value, "".__class__): # Single name.
            name = dropdown_value
        # elif len(dropdown_value) > 0: # Multilist, we only care about one item.
        #     name = dropdown_value[0]
        else:
            name = dropdown_value
        if isdummy(name):
            return refresh_all_data(username, token)
        if iswild(name):
            return refresh_all_data(username, token, indwild = True)
        return gr.Dropdown()
    
    # Dummy load only.
    multi_repo_input.change(
        fn=check_dummy,
        inputs=[hf_user, hf_token, multi_repo_input],
        outputs=[multi_repo_input]
    )
    
    multi_btn.click(
        fn=grant_access_multiple_repos,
        inputs=[hf_token, multi_repo_input, multi_user_input],
        outputs=multi_output
    )

if __name__ == "__main__":
    demo.launch()