from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import gradio as gr
import os
import sys
import random
import string
import time
from queue import Queue
from threading import Thread
import requests
import io
from PIL import Image
import base64
from deep_translator import GoogleTranslator

app = FastAPI()

API_URL = "https://api-inference.huggingface.co/models/jayavibhav/anime-dreamlike"
API_TOKEN = os.getenv("HF_READ_TOKEN") # it is free
headers = {"Authorization": f"Bearer {API_TOKEN}"}





text_gen = gr.Interface.load("models/Gustavosta/MagicPrompt-Stable-Diffusion")


queue = Queue()
queue_threshold = 100

def add_random_noise(prompt, noise_level=0.00):
    if noise_level == 0:
        noise_level = 0.00
    percentage_noise = noise_level * 5
    num_noise_chars = int(len(prompt) * (percentage_noise / 100))
    noise_indices = random.sample(range(len(prompt)), num_noise_chars)
    prompt_list = list(prompt)
    noise_chars = list(string.ascii_letters + string.punctuation + ' ' + string.digits)
    noise_chars.extend(['😍', '💩', '😂', '🤔', '😊', '🤗', '😭', '🙄', '😷', '🤯', '🤫', '🥴', '😴', '🤩', '🥳', '😔', '😩', '🤪', '😇', '🤢', '😈', '👹', '👻', '🤖', '👽', '💀', '🎃', '🎅', '🎄', '🎁', '🎂', '🎉', '🎈', '🎊', '🎮', '❤️', '💔', '💕', '💖', '💗', '🐶', '🐱', '🐭', '🐹', '🦊', '🐻', '🐨', '🐯', '🦁', '🐘', '🔥', '🌧️', '🌞', '🌈', '💥', '🌴', '🌊', '🌺', '🌻', '🌸', '🎨', '🌅', '🌌', '☁️', '⛈️', '❄️', '☀️', '🌤️', '⛅️', '🌥️', '🌦️', '🌧️', '🌩️', '🌨️', '🌫️', '☔️', '🌬️', '💨', '🌪️', '🌈'])
    for index in noise_indices:
        prompt_list[index] = random.choice(noise_chars)
    return "".join(prompt_list)

# Existing code...

import uuid  # Import the UUID library

# Existing code...

# Existing code...

request_counter = 0  # Global counter to track requests


    
def generate_image(inputs, is_negative, steps, cfg_scale, seed):
    try:
        global request_counter
        request_counter += 1
        timestamp = f"{time.time()}_{request_counter}"
        
        # Translate inputs to English
        translator_to_en = GoogleTranslator(source='auto', target='english')
        english_inputs = translator_to_en.translate(inputs)
        
        prompt_with_noise = add_random_noise(english_inputs) + f" - {timestamp}"
        payload = {
            "inputs": prompt_with_noise,
            "is_negative": is_negative,
            "steps": steps,
            "cfg_scale": cfg_scale,
            "seed": seed if seed is not None else random.randint(-1, 2147483647)
        }
        response = requests.post(API_URL, headers=headers, json=payload)
        response.raise_for_status()  # Raise an exception for HTTP errors
        image_bytes = response.content
        image = Image.open(io.BytesIO(image_bytes))
        return image
    except requests.exceptions.HTTPError as e:
        # Handle any HTTP errors
        print(f"HTTP Error: {e}")
        return "An unexpected error occurred while generating the image."
    except Exception as e:
        # Handle other exceptions
        print(f"Error generating image: {e}")
        return "An unexpected error occurred while generating the image. Please try again later."




def get_prompts(prompt_text):
    if not prompt_text:
        return "Please enter text before generating prompts." 
        raise gr.Error("Please enter text before generating prompts.")
    else:
        global request_counter
        request_counter += 1
        timestamp = f"{time.time()}_{request_counter}"

    options = [
        "photo anime, masterpiece, high quality, absurdres, "
        # Add other prompt options here...
    ]

    if prompt_text:
        chosen_option = random.choice(options)
        return text_gen(f"{prompt_text}, {chosen_option} - {timestamp}")
    else:
        return text_gen("", timestamp)

def initialize_api_connection():
    global headers
    API_TOKEN = os.getenv("HF_READ_TOKEN") # it is free
    headers = {"Authorization": f"Bearer {API_TOKEN}"}

# Run initialization functions on startup
initialize_api_connection()

@app.get("/generate_prompts")
def generate_prompts(prompt_text: str):
    return get_prompts(prompt_text)



from fastapi import Query

from fastapi import HTTPException

@app.get("/send_inputs")
def send_inputs(
    inputs: str,
    noise_level: float,
    is_negative: str,
    steps: int = 20,
    cfg_scale: int = 4.5,
    seed: int = None
):
    try:
        generated_image = generate_image(inputs, is_negative, steps, cfg_scale, seed)
        if generated_image is not None:
            image_bytes = io.BytesIO()
            generated_image.save(image_bytes, format="JPEG")
            image_base64 = base64.b64encode(image_bytes.getvalue()).decode("utf-8")
            return {"image_base64": image_base64}
        else:
            # Return an error message if the image couldn't be generated
            raise HTTPException(status_code=500, detail="Failed to generate image.")
    except Exception as e:
        # Log the error and return an error message
        print(f"Error generating image: {e}")
        raise HTTPException(status_code=500, detail="Failed to generate image.")


app.mount("/", StaticFiles(directory="static", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/static/index.html", media_type="text/html")