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

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

import spaces
from spaces.zero.client import _get_token

# Set up environment variables (you can also set these in your Spaces config)
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"] = "0.0.0.0"

# Create FastAPI app
app = FastAPI()

# Create and configure static directory
static_dir = Path("./static")
static_dir.mkdir(parents=True, exist_ok=True)

# Mount static directory to FastAPI
app.mount("/static", StaticFiles(directory="static"), name="static")

# Tell Gradio which paths are allowed to be served
os.environ["GRADIO_ALLOWED_PATHS"] = str(static_dir.resolve())

@spaces.GPU(duration=4*60)
def process_text(text):
    """Example GPU function - in reality, this might be model inference"""
    return text.upper()

def process_and_save(request: gr.Request, text):
    """Main processing function that handles tokens and calls GPU function"""
    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  # Simplified for demonstration

# Mark main function as not requiring GPU
process_and_save.zerogpu = True

# Create Gradio interface
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 Gradio app to FastAPI with SSR mode and node_port set to 7861
app = gr.mount_gradio_app(app, demo, path="/", ssr_mode=True, node_port=7861)

# Run the server on port 7860 (only one public port on Spaces)
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)