Update app.py

app.py

@@ -1,76 +1,237 @@
 import os
+import uuid
+import gradio as gr
+import requests
 from fastapi import FastAPI
 from fastapi.responses import FileResponse
-import gradio as gr
+from gradio_client import Client, handle_file
 
 # -------------------------
 # CONFIG
 # -------------------------
 
-FILE_DIR = "files"
-os.makedirs(FILE_DIR, exist_ok=True)
+SPACE = "pockit-cloud/main"
+client = Client(SPACE)
+
+HF_TOKEN = os.getenv("HF_TOKEN")  # Secret for private dataset
 
-# Your Space URL (auto-detect friendly)
-SPACE_NAME = os.getenv("SPACE_ID")
-if SPACE_NAME:
-    BASE_URL = f"https://{SPACE_NAME.split('/')[-1]}.hf.space"
-else:
-    BASE_URL = "http://127.0.0.1:7860"
+# -------------------------
+# LOGIN
+# -------------------------
+def login(user_id, password):
+    if not user_id or not password:
+        return "Username and password required", []
+
+    try:
+        files, status = client.predict(
+            user_id=user_id,
+            password=password,
+            api_name="/get_files_secure"
+        )
+        return status, files
+    except Exception as e:
+        return f"Login error: {e}", []
 
 # -------------------------
-# FASTAPI APP
+# LIST FILES
 # -------------------------
+def list_files(user_id, password):
+    if not user_id or not password:
+        return "Username and password required", []
+
+    try:
+        files, status = client.predict(
+            user_id=user_id,
+            password=password,
+            api_name="/get_files_secure"
+        )
+        return status, files
+    except Exception as e:
+        return str(e), []
 
-server = FastAPI()
+# -------------------------
+# DOWNLOAD (UNCHANGED LOGIC)
+# Returns temp file so Gradio generates link
+# -------------------------
+def download_file(user_id, password, filename):
 
-@server.get("/download/{filename}")
-def download_file(filename: str):
-    file_path = os.path.join(FILE_DIR, filename)
+    if not user_id or not password:
+        return None
 
-    if os.path.exists(file_path):
-        return FileResponse(
-            path=file_path,
-            filename=filename
+    if not filename:
+        return None
+
+    try:
+        link = client.predict(
+            user_id=user_id,
+            password=password,
+            filename=filename,
+            api_name="/get_download_link_action"
         )
 
-    return {"error": "File not found"}
+        headers = {}
+        if HF_TOKEN:
+            headers["Authorization"] = f"Bearer {HF_TOKEN}"
+
+        r = requests.get(link, headers=headers, stream=True)
+        r.raise_for_status()
+
+        temp_name = f"{uuid.uuid4()}_{filename}"
+
+        with open(temp_name, "wb") as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+
+        return temp_name  # IMPORTANT: keeps auto link behavior
+
+    except Exception as e:
+        print(e)
+        return None
+
+# -------------------------
+# UPLOAD
+# -------------------------
+def upload_file(user_id, password, file, custom_name):
+
+    if not user_id or not password:
+        return "Username and password required"
+
+    if not file:
+        return "No file selected"
+
+    try:
+        result = client.predict(
+            user_id=user_id,
+            password=password,
+            filepath=handle_file(file),
+            custom_name=custom_name if custom_name else os.path.basename(file),
+            api_name="/upload_file_secure"
+        )
+        return result[0]
+    except Exception as e:
+        return str(e)
 
 # -------------------------
-# LINK GENERATOR
+# DELETE
 # -------------------------
+def delete_file(user_id, password, filename):
+
+    if not user_id or not password:
+        return "Username and password required"
 
-def generate_link(filename):
     if not filename:
-        return "Enter filename"
+        return "No file selected"
+
+    try:
+        result = client.predict(
+            user_id=user_id,
+            password=password,
+            filename=filename,
+            api_name="/delete_file_secure"
+        )
+        return result
+    except Exception as e:
+        return str(e)
 
-    file_path = os.path.join(FILE_DIR, filename)
+# -------------------------
+# PASSWORD CHANGE
+# -------------------------
+def change_password(user_id, password, new_password):
 
-    if not os.path.exists(file_path):
-        return "File does not exist in /files"
+    if not user_id or not password or not new_password:
+        return "All fields required"
 
-    return f"{BASE_URL}/download/{filename}"
+    try:
+        result = client.predict(
+            user_id=user_id,
+            new_password=new_password,
+            api_name="/update_password"
+        )
+        return result
+    except Exception as e:
+        return str(e)
 
 # -------------------------
 # GRADIO UI
 # -------------------------
 
-with gr.Blocks() as app:
-    gr.Markdown("# 🔗 Direct Link Generator (HF Spaces)")
+with gr.Blocks(title="Pockit Cloud Client") as app:
+
+    gr.Markdown("# ☁️ Pockit Cloud Client")
+
+    user = gr.Textbox(label="User ID")
+    pw = gr.Textbox(label="Password", type="password")
+
+    with gr.Tab("Login"):
+        login_btn = gr.Button("Login")
+        login_status = gr.Textbox(label="Status")
+        login_files = gr.Dropdown(label="Files", choices=[])
 
-    filename_input = gr.Textbox(label="Filename")
-    link_output = gr.Textbox(label="Direct Link")
+        login_btn.click(
+            login,
+            inputs=[user, pw],
+            outputs=[login_status, login_files]
+        )
+
+    with gr.Tab("Files"):
+
+        refresh_btn = gr.Button("Refresh File List")
+        file_list = gr.Dropdown(label="Your Files", choices=[])
+        status_box = gr.Textbox(label="Status")
 
-    btn = gr.Button("Generate Link")
+        refresh_btn.click(
+            list_files,
+            inputs=[user, pw],
+            outputs=[status_box, file_list]
+        )
+
+        gr.Markdown("### Download")
 
-    btn.click(
-        generate_link,
-        inputs=filename_input,
-        outputs=link_output
-    )
+        download_btn = gr.Button("Download File")
+        download_output = gr.File(label="Downloaded File")
 
-# 🔥 Correct mounting (IMPORTANT FIX)
-app = gr.mount_gradio_app(server, app, path="/")
+        download_btn.click(
+            download_file,
+            inputs=[user, pw, file_list],
+            outputs=download_output
+        )
+
+        gr.Markdown("### Upload")
+
+        upload_input = gr.File(label="Select File")
+        custom_name = gr.Textbox(label="Custom Name (Optional)")
+        upload_btn = gr.Button("Upload")
+        upload_status = gr.Textbox(label="Upload Status")
+
+        upload_btn.click(
+            upload_file,
+            inputs=[user, pw, upload_input, custom_name],
+            outputs=upload_status
+        )
+
+        gr.Markdown("### Delete")
+
+        delete_btn = gr.Button("Delete Selected File")
+        delete_status = gr.Textbox(label="Delete Status")
+
+        delete_btn.click(
+            delete_file,
+            inputs=[user, pw, file_list],
+            outputs=delete_status
+        )
+
+    with gr.Tab("Password"):
+        new_pw = gr.Textbox(label="New Password", type="password")
+        change_btn = gr.Button("Change Password")
+        change_status = gr.Textbox(label="Status")
+
+        change_btn.click(
+            change_password,
+            inputs=[user, pw, new_pw],
+            outputs=change_status
+        )
 
-# Launch for Spaces
+# IMPORTANT FOR SPACES
 app.queue()
 app.launch()