###############################################################################
# 1) Set environment variables BEFORE importing Gradio (if you actually need them).
#    In many cases, you can omit these entirely. 
###############################################################################
import os

# If you want Gradio to run on a particular host/port, you can do this:
os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
os.environ["GRADIO_SERVER_PORT"] = "7860"

# If you do NOT really need GRADIO_ROOT_PATH, don’t set it. 
# If you do set it, do so BEFORE the Gradio import, e.g.:
os.environ["GRADIO_ROOT_PATH"] = "/_app/immutable"

###############################################################################
# 2) Now import everything
###############################################################################
import time
import json
import base64
from datetime import datetime
from pathlib import Path

import gradio as gr
from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
import uvicorn

# Hugging Face Spaces
import spaces
from spaces.zero.client import _get_token

###############################################################################
# 3) Create your FastAPI app and (optionally) mount a static folder for user files
###############################################################################
app = FastAPI()

static_dir = Path("./static")
static_dir.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")

###############################################################################
# 4) Define your GPU function and main processing function
###############################################################################
@spaces.GPU(duration=240)  # specify GPU usage for 4 minutes
def process_text(text):
    """Simulate a GPU-based process."""
    time.sleep(10)
    return text.upper()

def process_and_save(request: gr.Request, text: str):
    """Handles GPU call and writes result to a file in ./static."""
    token = _get_token(request)
    payload = token.split(".")[1]
    payload = f"{payload}{'=' * ((4 - len(payload) % 4) % 4)}"
    payload = json.loads(base64.urlsafe_b64decode(payload).decode())
    print(f"Token payload: {payload}")

    result = process_text(text)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    file_path = static_dir / f"output_{timestamp}.txt"
    with open(file_path, "w") as f:
        f.write(result)

    return gr.File(value=str(file_path))

# Mark as not requiring GPU
process_and_save.zerogpu = True

###############################################################################
# 5) Build the Gradio Blocks interface
###############################################################################
with gr.Blocks() as demo:
    text_input = gr.Textbox(label="Enter some text")
    submit_btn = gr.Button("Process and Download")
    output = gr.File(label="Download Processed File")

    submit_btn.click(
        fn=process_and_save,
        inputs=[text_input],
        outputs=output
    )

###############################################################################
# 6) Mount the Gradio app WITH SSR. Don’t manually mount _app/immutable.
###############################################################################
app = gr.mount_gradio_app(app, demo, path="/", ssr_mode=False)

###############################################################################
# 7) Run with Uvicorn
###############################################################################
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)