from flask import Flask, request, jsonify, send_file, render_template_string, make_response
import requests
import io
import os
import random
from PIL import Image
from deep_translator import GoogleTranslator

app = Flask(__name__)

API_URL = "https://api-inference.huggingface.co/models/Ojimi/anime-kawai-diffusion"
API_TOKEN = os.getenv("HF_READ_TOKEN")
headers = {"Authorization": f"Bearer {API_TOKEN}"}
timeout = 300  # タイムアウトを300秒に設定

# Function to query the API and return the generated image
def query(prompt, negative_prompt="", steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
    if not prompt:
        return None, "Prompt is required"

    key = random.randint(0, 999)
    
    # Translate the prompt from Russian to English if necessary
    prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
    print(f'Generation {key} translation: {prompt}')

    # Add some extra flair to the prompt
    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'Generation {key}: {prompt}')
    
    payload = {
        "inputs": prompt,
        "is_negative": False,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength,
        "parameters": {
            "width": width,
            "height": height
        }
    }

    for attempt in range(3):  # 最大3回の再試行
        try:
            response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
            if response.status_code != 200:
                return None, f"Error: Failed to get image. Status code: {response.status_code}, Details: {response.text}"

            image_bytes = response.content
            image = Image.open(io.BytesIO(image_bytes))
            return image, None
        except requests.exceptions.Timeout:
            if attempt < 2:  # 最後の試行でない場合は再試行
                print("Timeout occurred, retrying...")
                continue
            return None, "Error: The request timed out. Please try again."
        except requests.exceptions.RequestException as e:
            return None, f"Request Exception: {str(e)}"
        except Exception as e:
            return None, f"Error when trying to open the image: {e}"

# Content-Security-Policyヘッダーを設定するための関数
@app.after_request
def add_security_headers(response):
    response.headers['Content-Security-Policy'] = (
        "default-src 'self'; "
        "connect-src 'self' ^https?:\/\/[\w.-]+\.[\w.-]+(\/[\w.-]*)*(\?[^\s]*)?$"
        "img-src 'self' data:; "
        "style-src 'self' 'unsafe-inline'; "
        "script-src 'self' 'unsafe-inline'; "
    )
    return response

# HTML template for the index page
index_html = """

"""

@app.route('/')
def index():
    return render_template_string(index_html)

@app.route('/generate', methods=['GET'])
def generate_image():
    prompt = request.args.get("prompt", "")
    negative_prompt = request.args.get("negative_prompt", "")
    steps = int(request.args.get("steps", 35))
    cfg_scale = float(request.args.get("cfgs", 7))
    sampler = request.args.get("sampler", "DPM++ 2M Karras")
    strength = float(request.args.get("strength", 0.7))
    seed = int(request.args.get("seed", -1))
    width = int(request.args.get("width", 1024))
    height = int(request.args.get("height", 1024))

    image, error = query(prompt, negative_prompt, steps, cfg_scale, sampler, seed, strength, width, height)

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

    img_bytes = io.BytesIO()
    image.save(img_bytes, format='PNG')
    img_bytes.seek(0)
    return send_file(img_bytes, mimetype='image/png')

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