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

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

import spaces
from spaces.zero.client import _get_token

# Set environment variables
os.environ["GRADIO_SSR_MODE"] = "True"
os.environ["GRADIO_SERVER_PORT"] = "7860"
os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
os.environ["GRADIO_NODE_SERVER_NAME"] = "127.0.0.1"  # host only
os.environ["GRADIO_ROOT_PATH"] = "/"

# Create FastAPI app
app = FastAPI()

# Add middleware using the decorator
@app.middleware("http")
async def ssr_path_rewrite(request: Request, call_next):
    path = request.scope.get("path", "")
    if path.startswith("/7861"):
        remainder = path[len("/7861"):]
        if not remainder.startswith("/"):
            remainder = "/" + remainder
        new_path = "/_app" + remainder
        request.scope["path"] = new_path
        print(f"Rewrote path from {path} to {new_path}")
    response = await call_next(request)
    return response

# Set up static files (if needed)
static_dir = Path("./static")
static_dir.mkdir(parents=True, exist_ok=True)
app.mount("/static", StaticFiles(directory="static"), name="static")
os.environ["GRADIO_ALLOWED_PATHS"] = str(static_dir.resolve())

@spaces.GPU(duration=4*60)
def process_text(text):
    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)
    return processed_text

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.Textbox(label="Output")
    submit_btn.click(fn=process_and_save, inputs=[text_input], outputs=output)

# Mount the Gradio app with SSR mode, using node_port 7861
app = gr.mount_gradio_app(app, demo, path="/", ssr_mode=True, node_port=7861)

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