import gradio as gr
import requests
import io
import random
import os
import time
from PIL import Image
from deep_translator import GoogleTranslator
import json

# Project by Nymbo

API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
API_TOKEN = os.getenv("HF_READ_TOKEN")
headers = {"Authorization": f"Bearer {API_TOKEN}"}
timeout = 100

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

    key = random.randint(0, 999)
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    # Translate the prompt from Russian to English if necessary
    prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
    print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')
    # Add some extra flair to the prompt
    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'\033[1mGeneration {key}:\033[0m {prompt}')
    # Prepare the payload for the API call, including width and height
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength,
        "parameters": {
            "width": width,  # Pass the width to the API
            "height": height  # Pass the height to the API
        }
    }

    # Send the request to the API and handle the response
    response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
    if response.status_code != 200:
        print(f"Error: Failed to get image. Response status: {response.status_code}")
        print(f"Response content: {response.text}")
        if response.status_code == 503:
            raise gr.Error(f"{response.status_code} : The model is being loaded")
        raise gr.Error(f"{response.status_code}")
    try:
        # Convert the response content into an image
        image_bytes = response.content
        image = Image.open(io.BytesIO(image_bytes))
        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
        return image
    except Exception as e:
        print(f"Error when trying to open the image: {e}")
        return None

# CSS to style the app
css = """


body {
    background-color: #000;
    margin: 0;
    height: 100vh;
    display: flex;
    justify-content: center;
    align-items: center;
}

#neon-square {
    width: 300px;
    height: 300px;
    position: relative;
}

#neon-square::before {
    content: '';
    position: absolute;
    top: -20px;
    left: -20px;
    right: -20px;
    bottom: -20px;
    z-index: -1;
    filter: blur(20px);
}

#app {
    max-width: 1400px;
    position: relative; /* Allow positioning for neon effect */
}

.neon-cursor {
    position: absolute;
    top: 0;
    left: 0;
    width: 100%;
    height: 100%;
    pointer-events: none; /* Allow clicks through the neon effect */
}
#app-container {
    max-width: 1400px;
    margin-left: auto;
    margin-right: auto;
    overflow: hidden;
    touch-action: pan-up;
    color: #ff000f;
    font-family: 'Montserrat', sans-serif;
    text-align: center;
    text-shadow: 0 0 5px #1121, 0 0 20px #000, 0 0 30px #000;
}
.gradio-container {
    background: url(https://huggingface.co/spaces/K00B404/FLUX.1-Dev-Serverless-darn-enhanced-prompt/resolve/main/edge.png);
    background-size: 900px 880px;
    background-repeat: no-repeat;
    background-position: center;
    color:#000;
}
.dark\:bg-gray-950:is(.dark *) {
  --tw-bg-opacity: 1;
  background-color: rgb(157, 17, 142);
}
.gradio-container-4-41-0 .prose :last-child {
  margin-top: 8px !important;
}
.gradio-container-4-41-0 .prose :last-child {
  margin-bottom: -7px !important;
}
.dark {
    --button-primary-background-fill: #09e60d70;
    --button-primary-background-fill-hover: #00000070;
    --background-fill-primary: #000;
    --background-fill-secondary: #000;
}
.hide-container {
    margin-top;-2px;
}
#app-container3 {
    background-color: rgba(255, 255, 255, 0.001);  /* Corrected to make semi-transparent */
    max-width: 600px;
    margin-left: auto;
    margin-right: auto;
    margin-bottom: 10px;
    border-radius: 125px;
    box-shadow: 0 0 10px rgba(0,0,0,0.1); /* Adjusted shadow opacity */
}
#app-container {
    background-color: rgba(255, 255, 255, 0.001);  /* Semi-transparent background */
    max-width: 600px;
    margin: 0 auto;  /* Center horizontally */
    padding-bottom: 10px;
    border-radius: 25px;
    box-shadow: 0 0 10px rgba(0, 0, 0, 0.1); /* Adjusted shadow opacity */
}
#title-container {
    display: flex;
    align-items: center
    margin-bottom:10px;
    justify-content: center;
}
#title-icon {
    width: 32px;
    height: auto;
    margin-right: 10px;
}
#title-text {
    font-size: 30px;
    font-weight: bold;
    color: #000;
}

body, html, #app {
  margin: 0;
  width: 100%;
  height: 100%;
}

#app {
  overflow: hidden;
  touch-action: pan-up;
  color: #ff000f;
  font-family: 'Montserrat', sans-serif;
  text-align: center;
  text-shadow: 0 0 5px #1121, 0 0 20px #000, 0 0 30px #000;
}

#app h1 {
  --fontSize: 60px;
  --lineHeight: 80px;
  width: auto;
  height: calc(2 * var(--lineHeight));
  line-height: var(--lineHeight);
  margin: calc(50vh - var(--lineHeight)) auto 0;
  font-size: var(--fontSize);
  text-transform: uppercase;
}

#app a {
  margin-top: 10px;
  display: inline-block;
  text-decoration: none;
  color: #fff;
}

#app canvas {
  display: block;
  position: fixed;
  z-index: -1;
  top: 0;
}

"""

js2='''const neonSquare = document.getElementById('neon-square');
        let time = 0;

        function hslToRgb(h, s, l) {
            let r, g, b;
            if (s === 0) {
                r = g = b = l;
            } else {
                const hue2rgb = (p, q, t) => {
                    if (t < 0) t += 1;
                    if (t > 1) t -= 1;
                    if (t < 1/6) return p + (q - p) * 6 * t;
                    if (t < 1/2) return q;
                    if (t < 2/3) return p + (q - p) * (2/3 - t) * 6;
                    return p;
                };
                const q = l < 0.5 ? l * (1 + s) : l + s - l * s;
                const p = 2 * l - q;
                r = hue2rgb(p, q, h + 1/3);
                g = hue2rgb(p, q, h);
                b = hue2rgb(p, q, h - 1/3);
            }
            return [Math.round(r * 255), Math.round(g * 255), Math.round(b * 255)];
        }

        function updateColors() {
            const noise1 = noise.simplex2(time * 0.5, 0) * 0.5 + 0.5;
            const noise2 = noise.simplex2(time * 0.5, 1000) * 0.5 + 0.5;
            const noise3 = noise.simplex2(time * 0.5, 2000) * 0.5 + 0.5;
            const noise4 = noise.simplex2(time * 0.5, 3000) * 0.5 + 0.5;

            const [r1, g1, b1] = hslToRgb(noise1, 1, 0.5);
            const [r2, g2, b2] = hslToRgb(noise2, 1, 0.5);
            const [r3, g3, b3] = hslToRgb(noise3, 1, 0.5);
            const [r4, g4, b4] = hslToRgb(noise4, 1, 0.5);

            const gradient = `linear-gradient(45deg, 
                rgb(${r1},${g1},${b1}), 
                rgb(${r2},${g2},${b2}), 
                rgb(${r3},${g3},${b3}), 
                rgb(${r4},${g4},${b4}))`;
            
            neonSquare.style.background = gradient;
            neonSquare.style.boxShadow = `
                0 0 20px rgb(${r1},${g1},${b1}),
                0 0 40px rgb(${r2},${g2},${b2}),
                0 0 60px rgb(${r3},${g3},${b3}),
                0 0 80px rgb(${r4},${g4},${b4})
            `;
            
            neonSquare.style.setProperty('--glow', gradient);

            time += 0.01;
        }

        // Update colors every frame
        function animate() {
            updateColors();
            requestAnimationFrame(animate);
        }

        animate();'''

js_script="""
import { neonCursor } from 'https://unpkg.com/threejs-toys@0.0.8/build/threejs-toys.module.cdn.min.js';

let isNeonActive = false; // Toggle state for neon effect
let isFixedPathMode = false; // Toggle between edge-following and fixed path modes
let neonAnimationId = null; // To store the requestAnimationFrame ID

function moveNeonCursorAlongEdges() {
  const panel = document.getElementById('app');
  const rect = panel.getBoundingClientRect();

  const { top, left, right, bottom, width, height } = rect;

  let x, y;
  let time = performance.now() / 1000; // Time for smooth movement along edges
  let edgePosition = (time * 100) % (2 * width + 2 * height); // Total perimeter of the panel

  if (edgePosition < width) { // Top edge
    x = left + edgePosition;
    y = top;
  } else if (edgePosition < width + height) { // Right edge
    x = right;
    y = top + (edgePosition - width);
  } else if (edgePosition < 2 * width + height) { // Bottom edge
    x = right - (edgePosition - width - height);
    y = bottom;
  } else { // Left edge
    x = left;
    y = bottom - (edgePosition - 2 * width - height);
  }

  neonCursor({
    el: document.getElementById('neon-cursor'),
    shaderPoints: 16,
    curvePoints: 80,
    curveLerp: 0.5,
    radius1: 45,
    radius2: 30,
    velocityTreshold: 10,
    sleepRadiusX: 100,
    sleepRadiusY: 100,
    sleepTimeCoefX: 0.0025,
    sleepTimeCoefY: 0.0025,
    x, y // Position the neon cursor at the calculated x, y coordinates
  });

  if (isNeonActive && !isFixedPathMode) {
    neonAnimationId = requestAnimationFrame(moveNeonCursorAlongEdges);
  }
}

function toggleNeonEffect(event) {
  if (event.button === 2) { // Right-click
    isNeonActive = !isNeonActive;

    if (isNeonActive) {
      isFixedPathMode = !isFixedPathMode;

      if (isFixedPathMode) {
        moveNeonCursorInFixedPath(); // Start fixed path mode
      } else {
        moveNeonCursorAlongEdges(); // Start edge-following mode
      }
    } else {
      cancelAnimationFrame(neonAnimationId); // Stop the neon effect
    }
  }
}

window.addEventListener('contextmenu', (event) => {
  event.preventDefault();
  toggleNeonEffect(event);
});

window.addEventListener('resize', () => {
  if (isNeonActive) {
    cancelAnimationFrame(neonAnimationId);
    moveNeonCursorAlongEdges(); // Restart edge-following mode
  }
});
"""


title="""<center><h1>FluxCapacitor</h1></center>"""

# Build the Gradio UI with Blocks
with gr.Blocks( css=css, js=js2) as app: #theme='Nymbo/Nymbo_Theme',
    # Add a title to the app
    gr.HTML(title)

    # Add the neon cursor container
    with gr.Row():
        gr.HTML('<div id="neon-cursor" class="neon-cursor"></div>')
        
    # Container for all the UI elements
    with gr.Column(elem_id="app-container"):
        # Add a text input for the main prompt
        with gr.Row():
            with gr.Column(elem_id="prompt-container"):
                with gr.Row():
                    text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=2, elem_id="prompt-text-input")
                
                # Accordion for advanced settings
                with gr.Row():
                    with gr.Accordion("Advanced Settings", open=False):
                        negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="What should not be in the image", value="(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, misspellings, typos", lines=3, elem_id="negative-prompt-text-input")
                        with gr.Row():
                            width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=32)
                            height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1216, step=32)
                        steps = gr.Slider(label="Sampling steps", value=35, minimum=1, maximum=100, step=1)
                        cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=1)
                        strength = gr.Slider(label="Strength", value=0.7, minimum=0, maximum=1, step=0.001)
                        seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1) # Setting the seed to -1 will make it random
                        method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])

        # Add a button to trigger the image generation
        with gr.Row():
            text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
        
        # Image output area to display the generated image
        with gr.Row():
            image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
        
        # Bind the button to the query function with the added width and height inputs
        text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height], outputs=image_output)

# Launch the Gradio app
app.launch(show_api=False, share=False)