import os
import time
import requests
import gradio as gr
import uuid

# Hugging Face API
HF_API_KEY = os.getenv("HF_TOKEN")

# Chevereto API
CHEVERETO_API_URL = os.getenv("API_URL")
CHEVERETO_API_KEY = os.getenv("API_KEY")
CHEVERETO_ALBUM_ID = os.getenv("ALBUM_ID")
EXPECTED_API_KEY = os.getenv("API_PASS")

HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}

def generate_random_filename(extension: str) -> str:
    return str(uuid.uuid4()) + extension

def validate_api_key(api_key: str) -> bool:
    return api_key == EXPECTED_API_KEY

def generate_image(prompt: str, model_name: str):
    hf_api_url = f"https://api-inference.huggingface.co/models/{model_name}"
    response = requests.post(hf_api_url, headers=HEADERS, json={"inputs": prompt})
    
    if response.status_code == 200:
        image_data = response.content
        random_filename = generate_random_filename(".png")
        image_path = f"/tmp/{random_filename}"

        with open(image_path, "wb") as f:
            f.write(image_data)

        time.sleep(10)  # 10-second delay
        return image_path
    else:
        return f"Error generating image: {response.text}"

def upload_to_chevereto(image_path: str):
    with open(image_path, "rb") as img_file:
        files = {"source": img_file}
        data = {"key": CHEVERETO_API_KEY, "format": "json", "album_id": CHEVERETO_ALBUM_ID, "expiration": "PT10M"}
        
        response = requests.post(CHEVERETO_API_URL, files=files, data=data)
    
    if response.status_code == 200:
        return response.json().get("image", {}).get("url")
    else:
        return f"Error uploading image: {response.text}"

def generate_and_upload(prompt: str, model_name: str, api_key: str):
    if not validate_api_key(api_key):
        return "Invalid or missing API key", ""
    
    img_path = generate_image(prompt, model_name)
    if "Error" in img_path:
        return img_path, ""
    
    url = upload_to_chevereto(img_path)
    return img_path, url

# Gradio UI
with gr.Blocks(theme='NoCrypt/miku') as app:
    gr.HTML(
        """
        
        <h1 style='text-align: center; color: red; font-size: 20px; font-weight: bold;'>♥ Tasia API on Marty Space ♥</h1>
        
        """
    )

    with gr.Row():
        prompt_input = gr.Textbox(label="Enter a prompt:", placeholder="A futuristic spaceship...")
        model_input = gr.Textbox(label="Enter model name:", placeholder="stabilityai/stable-diffusion-3.5-large")
        api_key_input = gr.Textbox(label="Enter your API key:", type="password", placeholder="Your secret API key")
        generate_button = gr.Button("Generate & Upload")
    
    image_output = gr.Image(label="Generated Image")
    url_output = gr.Textbox(label="Image URL", interactive=False)

    generate_button.click(
        fn=generate_and_upload,
        inputs=[prompt_input, model_input, api_key_input],
        outputs=[image_output, url_output]
    )

    # Footer
    gr.HTML(
        """
        
        <h1 style='text-align: center; color: red; font-size: 20px; font-weight: bold;'>♥ Made by Tasia with ♥</h1>
        
        """
    )

app.launch(server_port=7860)