import os
import gradio as gr
from huggingface_hub import HfApi, CommitOperationAdd, create_commit
import pandas as pd
from datetime import datetime
import tempfile
from typing import Optional, Tuple

# Initialize Hugging Face client
hf_api = HfApi(token=os.getenv("HF_TOKEN"))
DATASET_REPO = "HuggingFaceTB/smolvlm2-iphone-waitlist"
MAX_RETRIES = 3

def commit_signup(username: str) -> Optional[str]:
    """Attempt to commit new signup atomically, always reading latest data before commit"""
    try:
        # Always get the latest data right before committing
        try:
            file_content = hf_api.hf_hub_download(
                repo_id=DATASET_REPO,
                repo_type="dataset",
                filename="waitlist.csv",
                token=os.getenv("HF_TOKEN")
            )
            current_data = pd.read_csv(file_content)
            if 'userid' not in current_data.columns or 'joined_at' not in current_data.columns:
                current_data = pd.DataFrame(columns=['userid', 'joined_at'])
        except Exception as e:
            current_data = pd.DataFrame(columns=['userid', 'joined_at'])
        # If user already exists in the latest data, don't add them again
        if username in current_data['userid'].values:
            return "exist"

        # Add new user with timestamp
        new_row = pd.DataFrame([{
            'userid': username,
            'joined_at': datetime.utcnow().isoformat()
        }])
        
        # Combine with latest data
        updated_data = pd.concat([current_data, new_row], ignore_index=True)

        # Save to temp file
        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as tmp:
            updated_data.to_csv(tmp.name, index=False)
            
            operation = CommitOperationAdd(
                path_in_repo="waitlist.csv",
                path_or_fileobj=tmp.name
            )

            try:
                create_commit(                    
                    repo_id=DATASET_REPO,
                    repo_type="dataset",
                    operations=[operation],
                    commit_message=f"Add user {username} to waitlist",
                    token=os.getenv("HF_TOKEN"),
                )
                os.unlink(tmp.name)
                return None
            except Exception as e:
                os.unlink(tmp.name)
                return str(e)
    except Exception as e:
        return str(e)

def join_waitlist(
    profile: gr.OAuthProfile | None,
    oauth_token: gr.OAuthToken | None
) -> gr.update:
    """Add user to the SmolVLM2 iPhone waitlist with retry logic for concurrent updates"""
    
    if profile is None or oauth_token is None:
        return gr.update(value="Please log in with your Hugging Face account first!")

    username = profile.username
    
    for attempt in range(MAX_RETRIES):
        try:
            # Try to commit the update (which will get latest data internally)
            error = commit_signup(username)
            
            if error == "exist":
                return gr.update(
                    value="## You are already in our waitlist, we will come back with news soon!",
                    visible=True
                )
            if error is None:  # Success or already registered
                # Verify user is in the list
                file_content = hf_api.hf_hub_download(
                    repo_id=DATASET_REPO,
                    repo_type="dataset",
                    filename="waitlist.csv",
                    token=os.getenv("HF_TOKEN")
                )
                verification_data = pd.read_csv(file_content)
                if username in verification_data['userid'].values:
                    return gr.update(
                        value="## 🎉 Thanks for joining the waitlist! We'll keep you updated on SmolVLM2 iPhone app.",
                        visible=True
                    )
                
                # If verification failed and we have retries left, try again
                if attempt < MAX_RETRIES - 1:
                    continue
            
            # If we got a conflict error, retry
            if error and "Conflict" in str(error) and attempt < MAX_RETRIES - 1:
                continue
                
            # Other error or last attempt
            if error:
                gr.Error(f"An error occurred: {str(error)}")
            return gr.update(
                value="## 🫠 Sorry, something went wrong. Please try again later.",
                visible=True
            )

        except Exception as e:
            print(str(e))
            if attempt == MAX_RETRIES - 1:  # Last attempt
                gr.Error(f"An error occurred: {str(e)}")
                return gr.update(
                    value="## 🫠 Sorry, something went wrong. Please try again later.",
                    visible=True
                )
                
def update_ui(profile: gr.OAuthProfile | None) -> Tuple[gr.update, gr.update, gr.update]:
    """Update UI elements based on login status"""
    if profile is None:
        return (
            gr.update(
                value="## Please sign in with your Hugging Face account to join the waitlist.",
                visible=True
            ),
            gr.update(visible=False),
            gr.update(visible=False)
        )
    return (
        gr.update(
            value=f"## Welcome {profile.name} 👋 Click the button below 👇 to receive any updates related with the SmolVLM2 iPhone application",
            visible=True
        ),
        gr.update(visible=True),
        gr.update(visible=True)
    )
# Create the interface
with gr.Blocks(title="SmolVLM2 Waitlist") as demo:
    gr.Markdown("""
    # Join SmolVLM2 iPhone Waitlist 📱
    ## Please log-in using your Hugging Face Id to continue
    """)

    gr.LoginButton()
    welcome_msg = gr.Markdown(visible=True)
    join_button = gr.Button("Join Waitlist", variant="primary", visible=False)
    status_msg = gr.Markdown(visible=False)
    
    # Update UI elements on load and auth change
    demo.load(
        fn=update_ui,
        inputs=None,
        outputs=[welcome_msg, join_button, status_msg]
    )
    
    # Handle join button click
    join_button.click(join_waitlist, None, status_msg)

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