import os
from pathlib import Path
import json, base64, time
from datetime import datetime

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

import spaces
from spaces.zero.client import _get_token

app = FastAPI()

# Create and mount a static directory if needed
static_dir = Path("./static")
static_dir.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")

# Remove or override the default GRADIO_ROOT_PATH
os.environ["GRADIO_ROOT_PATH"] = "/assets"

@spaces.GPU(duration=4*60)
def process_text(text):
    time.sleep(10)
    return text.upper()

def process_and_save(request: gr.Request, text):
    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}")

    processed_text = 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(processed_text)
    return gr.File(value=str(file_path))

process_and_save.zerogpu = True

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)

# Mount Gradio app with SSR
app = gr.mount_gradio_app(app, demo, path="/", ssr_mode=True)

# Make sure assets are served from /assets
assets_dir = Path("./assets")
assets_dir.mkdir(parents=True, exist_ok=True)
app.mount("/assets", StaticFiles(directory=str(assets_dir)), name="assets")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)