import streamlit as st
from huggingface_hub import HfApi, create_repo, upload_file
import os

# App title
st.title("HF Helper - Manage Your Hugging Face Spaces Easily")

# Use session state to store the token
if "hf_token" not in st.session_state:
    st.session_state["hf_token"] = None

# Ask for token input if not already saved
if st.session_state["hf_token"] is None:
    token_input = st.text_input("Enter your Hugging Face Access Token:", type="password")
    if token_input:
        st.session_state["hf_token"] = token_input
        st.write("Token saved.")
else:
    st.write("Token already saved.")

# Ensure token is valid before proceeding
if st.session_state["hf_token"] is None:
    st.error("Please provide a valid token to proceed.")
    st.stop()

# User Input for Space Name
space_name = st.text_input("Enter the name of the Hugging Face Space to update/create:")

# Code sections for `app.py` and `requirements.txt`
st.subheader("Provide your `app.py` Code")
app_code = st.text_area("Paste the code for `app.py` here:", height=200)

st.subheader("Provide your `requirements.txt` Code")
req_code = st.text_area("Paste the content for `requirements.txt` here:", height=100)

# Button to commit changes
if st.button("Commit Changes"):
    if not st.session_state["hf_token"].strip():
        st.error("Access token is required. Please enter a valid token.")
    elif not space_name.strip():
        st.error("Please enter a valid Space name.")
    elif not app_code.strip() or not req_code.strip():
        st.error("Both `app.py` and `requirements.txt` content are required.")
    else:
        # Initialize Hugging Face API
        api = HfApi()

        # Authenticate using the token
        try:
            token = st.session_state["hf_token"]
            user_info = api.whoami(token=token)
            st.success(f"Authenticated as {user_info['name']}")
        except Exception as e:
            st.error(f"Authentication failed: {e}")
            st.stop()

        try:
            # Full Space name
            space_full_name = f"{user_info['name']}/{space_name}"

            # Check if Space exists or create a new one
            try:
                st.info(f"Checking if Space '{space_full_name}' exists...")
                api.repo_info(repo_id=space_full_name, repo_type="space", token=token)
                st.success(f"Space '{space_full_name}' exists. Updating files...")
            except:
                st.info(f"Space '{space_full_name}' does not exist. Creating new Space...")
                create_repo(repo_id=space_name, repo_type="space", space_sdk="streamlit", token=token)

            # Upload `app.py`
            with open("app.py", "w") as app_file:
                app_file.write(app_code)
            upload_file(
                path_or_fileobj="app.py",
                path_in_repo="app.py",
                repo_id=space_full_name,
                repo_type="space",
                token=token
            )

            # Upload `requirements.txt`
            with open("requirements.txt", "w") as req_file:
                req_file.write(req_code)
            upload_file(
                path_or_fileobj="requirements.txt",
                path_in_repo="requirements.txt",
                repo_id=space_full_name,
                repo_type="space",
                token=token
            )

            st.success(f"Files successfully updated in Space '{space_full_name}'!")

        except Exception as e:
            st.error(f"An error occurred: {e}")
        finally:
            # Clean up temporary files
            if os.path.exists("app.py"):
                os.remove("app.py")
            if os.path.exists("requirements.txt"):
                os.remove("requirements.txt")