import gradio as gr
from huggingface_hub import login, logout, whoami

# Function to handle login
def handle_login(token):
    try:
        # Attempt to log in with the provided token
        login(token=token, add_to_git_credential=False)
        user_info = whoami()
        return f"Logged in as: {user_info['name']}"
    except Exception as e:
        # Handle login failure
        logout()  # Ensure the user is logged out if login fails
        return f"Login failed: {str(e)}"

# Function to check if the user is logged in
def is_logged_in():
    try:
        # Check if the user is authenticated
        whoami()
        return True
    except:
        return False

# Function to restrict access to the app
def restricted_functionality(prompt):
    if not is_logged_in():
        return "Please log in to use this feature."
    # Simulate model response (replace with actual model inference)
    return f"Model response to: {prompt}"

# Gradio interface
with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.Markdown("# Inference Provider")
        gr.Markdown("This Space showcases the deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct model, served by the nebius API. Sign in with your Hugging Face account to use this API.")
        token_input = gr.Textbox(label="Hugging Face Token", type="password")
        login_button = gr.Button("Sign in")
        login_status = gr.Markdown("")

    # Main app functionality
    with gr.Column(visible=False) as main_interface:  # Hide until logged in
        prompt = gr.Textbox(label="Your Prompt")
        output = gr.Textbox(label="Model Response")
        generate_button = gr.Button("Generate")

    # Load the model
    model_interface = gr.load("models/deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct", provider="nebius")

    # Handle login
    def update_interface(token):
        login_result = handle_login(token)
        if "Logged in as:" in login_result:
            return {main_interface: gr.update(visible=True), login_status: login_result}
        else:
            return {main_interface: gr.update(visible=False), login_status: login_result}

    login_button.click(update_interface, inputs=token_input, outputs=[main_interface, login_status])

    # Handle text generation (restricted to logged-in users)
    generate_button.click(restricted_functionality, inputs=prompt, outputs=output)

# Launch the app
demo.launch()