app.py

import os
import uuid
import gradio as gr
import requests
from fastapi import FastAPI
from fastapi.responses import FileResponse
from gradio_client import Client, handle_file

# -------------------------
# CONFIG
# -------------------------

SPACE = "pockit-cloud/main"
client = Client(SPACE)

HF_TOKEN = os.getenv("HF_TOKEN")  # Secret for private dataset

# -------------------------
# 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}", []

# -------------------------
# 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), []

# -------------------------
# DOWNLOAD (UNCHANGED LOGIC)
# Returns temp file so Gradio generates link
# -------------------------
def download_file(user_id, password, filename):

    if not user_id or not password:
        return None

    if not filename:
        return None

    try:
        link = client.predict(
            user_id=user_id,
            password=password,
            filename=filename,
            api_name="/get_download_link_action"
        )

        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)

# -------------------------
# DELETE
# -------------------------
def delete_file(user_id, password, filename):

    if not user_id or not password:
        return "Username and password required"

    if not 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)

# -------------------------
# PASSWORD CHANGE
# -------------------------
def change_password(user_id, password, new_password):

    if not user_id or not password or not new_password:
        return "All fields required"

    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(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=[])

        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")

        refresh_btn.click(
            list_files,
            inputs=[user, pw],
            outputs=[status_box, file_list]
        )

        gr.Markdown("### Download")

        download_btn = gr.Button("Download File")
        download_output = gr.File(label="Downloaded File")

        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
        )

# IMPORTANT FOR SPACES
app.queue()
app.launch()