import gradio as gr
import random
import os
from PIL import Image
from typing import Optional, Dict, Tuple
from huggingface_hub import InferenceClient

# Project by Nymbo

API_TOKEN = os.getenv("HF_READ_TOKEN")
timeout = 100

def flux_krea_generate(
    prompt: str,
    negative_prompt: str = "(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",
    steps: int = 35,
    cfg_scale: float = 7.0,
    sampler: str = "DPM++ 2M Karras",
    seed: int = -1,
    strength: float = 0.7,
    width: int = 1024,
    height: int = 1024,
) -> Optional[Image.Image]:
    """Generate a single image from a text prompt using FLUX.1-Krea-dev.

    Contract (for UI):
    - Inputs: prompt (required), optional tuning params below. No input image used.
    - Output: a PIL.Image for Gradio image component wiring.
    - Errors: raises gr.Error on auth/model/service issues.

    Args:
        prompt: Required. Describes what to create. Keep it specific and concise.
        negative_prompt: Phrases/objects to avoid in the output.
        steps: Number of diffusion steps (1-100). Higher may improve detail but is slower.
        cfg_scale: Classifier-free guidance scale (1-20). Higher forces closer adherence to prompt.
        sampler: Sampler algorithm label (UI only; provider may ignore or auto-select).
        seed: Set a deterministic seed (>=0). Use -1 to randomize per call.
        strength: Kept for parity; has no effect without an input image (0-1).
        width: Output width in pixels (64-1216). Prefer multiples of 32.
        height: Output height in pixels (64-1216). Prefer multiples of 32.

    Returns:
        PIL.Image or None if prompt is empty.
    """
    if prompt == "" or prompt is None:
        return None

    key = random.randint(0, 999)
    
    # Add some extra flair to the prompt
    enhanced_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'\033[1mGeneration {key}:\033[0m {enhanced_prompt}')
    
    try:
        # Initialize the Hugging Face Inference Client
        # Try different providers in order of preference
        providers = ["auto", "replicate", "fal-ai"]
        
        for provider in providers:
            try:
                client = InferenceClient(
                    api_key=API_TOKEN,
                    provider=provider
                )
                
                # Generate the image using the proper client
                image = client.text_to_image(
                    prompt=enhanced_prompt,
                    negative_prompt=negative_prompt,
                    model="black-forest-labs/FLUX.1-Krea-dev",
                    width=width,
                    height=height,
                    num_inference_steps=steps,
                    guidance_scale=cfg_scale,
                    seed=seed if seed != -1 else random.randint(1, 1000000000)
                )
                
                print(f'\033[1mGeneration {key} completed with {provider}!\033[0m ({enhanced_prompt})')
                return image
                
            except Exception as provider_error:
                print(f"Provider {provider} failed: {provider_error}")
                if provider == providers[-1]:  # Last provider
                    raise provider_error
                continue
        
    except Exception as e:
        print(f"Error during image generation: {e}")
        if "404" in str(e):
            raise gr.Error("Model not found. Please ensure the FLUX.1-Krea-dev model is accessible with your API token.")
        elif "503" in str(e):
            raise gr.Error("The model is currently being loaded. Please try again in a moment.")
        elif "401" in str(e) or "403" in str(e):
            raise gr.Error("Authentication failed. Please check your HF_READ_TOKEN environment variable.")
        else:
            raise gr.Error(f"Image generation failed: {str(e)}")
        return None


def _space_base_url() -> Optional[str]:
    """Return the public base URL of this Space if available.

    Looks up SPACE_ID (e.g. "username/space-name") and converts it to
    the public subdomain "https://username-space-name.hf.space".
    If SPACE_ID is not set, optionally respects HF_SPACE_BASE_URL.
    """
    # Explicit override if provided
    explicit = os.getenv("HF_SPACE_BASE_URL")
    if explicit:
        return explicit.rstrip("/")

    space_id = os.getenv("SPACE_ID")
    if not space_id:
        return None
    sub = space_id.replace("/", "-")
    return f"https://{sub}.hf.space"


def _save_image_and_url(image: Image.Image, key: int) -> Tuple[str, Optional[str]]:
    """Save the image to a temporary path and construct a public URL if running on Spaces.

    Returns (local_path, public_url_or_None).
    """
    # Ensure POSIX-like temp dir (Spaces is Linux). Still works locally.
    out_dir = os.path.join("/tmp", "flux-krea-outputs")
    os.makedirs(out_dir, exist_ok=True)
    file_path = os.path.join(out_dir, f"flux_{key}.png")
    image.save(file_path)

    base_url = _space_base_url()
    public_url = None
    if base_url:
        # Gradio serves local files via /file=<absolute-path>
        # Normalize backslashes to forward slashes in case of local dev on Windows
        posix_path = file_path.replace("\\", "/")
        public_url = f"{base_url}/file={posix_path}"

    return file_path, public_url


def flux_krea_generate_mcp(
    prompt: str,
    negative_prompt: str = "(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",
    steps: int = 35,
    cfg_scale: float = 7,
    sampler: str = "DPM++ 2M Karras",
    seed: int = -1,
    strength: float = 0.7,
    width: int = 1024,
    height: int = 1024,
) -> Dict[str, object]:
    """Generate an image (MCP tool) and return a JSON payload with a public URL.

    This endpoint is tailored for Model Context Protocol (MCP) clients per the
    latest Hugging Face MCP Space guidance (see hf-docs-search: "Spaces as MCP servers").

    Inputs:
      - prompt (str, required): Description of the desired image.
      - negative_prompt (str): Items to avoid in the generation.
      - steps (int, 1-100): Denoising steps. Higher is slower and may add detail.
      - cfg_scale (float, 1-20): Guidance strength. Higher adheres more to the prompt.
      - sampler (str): Sampler label (informational; provider may auto-select).
      - seed (int): -1 for random per call; otherwise a deterministic seed >= 0.
      - strength (float, 0-1): No-op in pure text-to-image; kept for cross-app parity.
      - width (int, 64-1216): Output width (prefer multiples of 32).
      - height (int, 64-1216): Output height (prefer multiples of 32).

    Returns (JSON):
      - image_path (str): Absolute path to the saved image on the Space VM.
      - image_url (str|None): Publicly accessible URL to the image on the Space
        (present when running on Spaces; None when running locally).
      - seed (int): The seed used for this run (randomized if input was -1).
      - width (int), height (int): Echo of output dimensions.
      - sampler (str): Echo of requested sampler.

    Error Modes:
      - Raises a Gradio-friendly error with a concise message for common HTTP
        failure codes (401/403 auth; 404 model; 503 warmup).
    """
    if not prompt:
        raise gr.Error("'prompt' is required and cannot be empty.")

    # Reuse core generator for image creation
    image = flux_krea_generate(
        prompt=prompt,
        negative_prompt=negative_prompt,
        steps=steps,
        cfg_scale=cfg_scale,
        sampler=sampler,
        seed=seed,
        strength=strength,
        width=width,
        height=height,
    )

    if image is None:
        raise gr.Error("No image generated.")

    # Save and build URLs
    key = random.randint(0, 999)
    file_path, public_url = _save_image_and_url(image, key)

    # Expose URL explicitly for MCP clients (LLMs need a resolvable URL)
    return {
        "image_path": file_path,
        "image_url": public_url,
        "seed": seed,
        "width": width,
        "height": height,
        "sampler": sampler,
    }

# CSS to style the app
css = """
#app-container {
    max-width: 800px;
    margin-left: auto;
    margin-right: auto;
}
"""

# Build the Gradio UI with Blocks
with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
    # Add a title to the app
    gr.HTML("<center><h1>FLUX.1-Krea-dev</h1></center>")
    gr.HTML("<center><p>High-quality image generation via Model Context Protocol</p></center>")
    
    # 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():
            # Output component only; no input image is required by the tool
            image_output = gr.Image(label="Image Output", elem_id="gallery")
        
        # Bind the button to the flux_krea_generate function for the UI only
        # Hide this event as an MCP tool to avoid schema confusion (UI wires image output)
        text_button.click(
            flux_krea_generate,
            inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height],
            outputs=image_output,
            show_api=False,
            api_description=False,
        )

    # Expose a dedicated MCP/API endpoint with a clear schema (text-to-image only)
    # This wrapper returns both a local file path and a fully-qualified public URL
    # when running on Spaces so LLMs can access the finished image.
    gr.api(
        flux_krea_generate_mcp,
        api_name="generate_image",
        api_description=(
            "Generate an image from a text prompt using FLUX.1-Krea-dev. "
            "Returns JSON with image_path and image_url (public URL when on Spaces)."
        ),
    )

# Launch the Gradio app with MCP server enabled
app.launch(show_api=True, share=False, mcp_server=True)