###############################################################################
# 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

from io import BytesIO
from fastapi import FastAPI, Form
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import StreamingResponse, FileResponse


# 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")

@app.post("/echo-image/",
          tags=["Test"],
          description="Descripción",
          summary="Summary"
          )
async def echo_image(image: UploadFile = File(...)):
    if not image.content_type.startswith("image/"):
        return {"error": "El archivo no es una imagen"}

    contents = await image.read()
    return StreamingResponse(BytesIO(contents), media_type=image.content_type)

###############################################################################
# 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)