import os
import requests
from flask import Flask, request, jsonify, render_template

app = Flask(__name__)

# Hugging Face API Configuration
HF_API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
HF_API_KEY = os.environ.get("HF_TOKEN")  # Set your Hugging Face API key

# Chevereto Configuration
CHEVERETO_API_URL = os.environ.get("API_URL")
CHEVERETO_API_KEY = os.environ.get("API_KEY")
CHEVERETO_ALBUM_ID = os.environ.get("ALBUM_ID")

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

def generate_image_and_upload(prompt: str) -> str:
    """Generates an image using Hugging Face API and uploads it to Chevereto."""
    
    response = requests.post(HF_API_URL, headers=HEADERS, json={"inputs": prompt})

    if response.status_code == 200:
        image_data = response.content
        image_path = "static/generated_image.png"
        
        with open(image_path, "wb") as f:
            f.write(image_data)

        return upload_to_chevereto(image_path)  # Upload and return URL
    else:
        return f"Error generating image: {response.text}"

def upload_to_chevereto(img_path: str) -> str:
    """Uploads an image to Chevereto and returns the public image URL."""
    
    with open(img_path, "rb") as image_file:
        files = {"source": image_file}
        data = {
            "key": CHEVERETO_API_KEY,
            "format": "json",
            "album": CHEVERETO_ALBUM_ID
        }
        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}"

@app.route("/", methods=["GET"])
def index():
    """Renders the Star Trek-themed UI."""
    return render_template("index.html")

@app.route("/generate", methods=["POST"])
def generate():
    """API Endpoint: Generate an image from a text prompt and upload it."""
    data = request.get_json()
    prompt = data.get("prompt", "")

    if not prompt:
        return jsonify({"error": "Missing prompt"}), 400

    image_url = generate_image_and_upload(prompt)
    return jsonify({"image_url": image_url})

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