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


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

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

# 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)  # Specify GPU duration in seconds, added for testing
def process_text(text):
    """Example GPU function - in reality, this might be model inference"""
    time.sleep(10)

    return text.upper()
# I've tried also wrapping the immediate function called by gradio, and the zerogpu = True trick

def process_and_save(request: gr.Request, text):
    """Main processing function that handles tokens and calls GPU function"""
    # Get and decode the authentication token
    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}")  # For debugging

    # Process the text using GPU function
    processed_text = process_text(text)
    
    # Save to file
    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)
    
    # Convert file_path to a string before returning
    return gr.File(value=str(file_path))

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

# Mount Gradio app to FastAPI with SSR mode for Spaces
app = gr.mount_gradio_app(app, demo, path="/", ssr_mode=True)

# Run server
if __name__ == "__main__":
    # Set SSR mode for Spaces
    os.environ["GRADIO_SSR_MODE"] = "True"
    uvicorn.run(app, host="0.0.0.0", port=7860)