import os
import gradio as gr

try:
    from huggingface_hub import InferenceClient
except ImportError:
    InferenceClient = None

STYLES = [
    "Cinematic",
    "Photorealistic",
    "Digital Art",
    "Fantasy",
    "Anime",
    "3D Render",
    "Editorial Photography",
    "Concept Art",
    "Minimalist",
    "Oil Painting",
]

def local_prompt(idea, style, aspect, mood, detail, camera):
    idea = idea.strip()
    style = style or "Cinematic"
    mood = mood or "Atmospheric"
    detail = detail or "highly detailed"
    camera = camera or "natural perspective"

    return f"""Create a {style.lower()} image of {idea}.

Subject and scene: Clearly depict the central subject with believable anatomy, proportions, materials, and environmental context. Build a coherent setting around the subject with foreground, midground, and background elements that support the story.

Composition: {camera}, strong visual hierarchy, intentional framing, balanced negative space, convincing depth and perspective.

Lighting and atmosphere: {mood.lower()} lighting, natural light falloff, dimensional shadows, subtle atmospheric depth, realistic highlights and reflections where appropriate.

Visual treatment: {detail}, rich textures, refined edges, nuanced surfaces, coherent color relationships, cinematic visual storytelling.

Quality: high visual fidelity, crisp focal detail, realistic spatial relationships, polished professional finish, no accidental text, no watermarks.

Aspect ratio: {aspect}.
"""

def generate(idea, style, aspect, mood, detail, camera, use_ai):
    if not idea or not idea.strip():
        return "Enter an image idea first."

    token = os.getenv("HF_TOKEN")
    model = os.getenv("PROMPT_MODEL", "Qwen/Qwen2.5-7B-Instruct")

    if use_ai and token and InferenceClient:
        try:
            client = InferenceClient(model=model, token=token)
            system = """You are an expert image-prompt architect. Expand a user's idea into one detailed,
production-ready prompt for an image-generation model. Preserve the user's core intent.
Do not invent major subjects. Include subject, environment, composition, lighting, style,
materials/textures, mood, camera/lens language when useful, quality, and aspect ratio.
Return only the finished prompt, with no preamble or quotation marks."""
            user = f"""Idea: {idea}
Style: {style}
Aspect ratio: {aspect}
Mood: {mood}
Detail level: {detail}
Camera/composition: {camera}"""
            result = client.chat_completion(
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                max_tokens=900,
                temperature=0.7,
            )
            return result.choices[0].message.content.strip()
        except Exception as e:
            return local_prompt(idea, style, aspect, mood, detail, camera) + \
                   f"\n\n[AI enhancement unavailable; used built-in generator. Error: {type(e).__name__}]"

    return local_prompt(idea, style, aspect, mood, detail, camera)

with gr.Blocks(title="Free Image Prompt Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """# ✦ Free Image Prompt Generator
### Turn a simple idea into a detailed, descriptive image prompt.

Describe what you want to create. The generator expands it into a structured prompt suitable for image-generation models."""
    )

    with gr.Row():
        with gr.Column(scale=2):
            idea = gr.Textbox(
                label="Your idea",
                placeholder="e.g. A lone astronaut exploring an abandoned underwater city",
                lines=5,
            )
            with gr.Row():
                style = gr.Dropdown(STYLES, value="Cinematic", label="Style")
                aspect = gr.Dropdown(
                    ["1:1", "4:5", "3:4", "16:9", "9:16", "21:9"],
                    value="16:9",
                    label="Aspect ratio",
                )
            with gr.Row():
                mood = gr.Textbox(value="Atmospheric", label="Mood / atmosphere")
                detail = gr.Textbox(value="highly detailed", label="Detail level")
            camera = gr.Textbox(
                value="medium-wide cinematic perspective",
                label="Camera / composition",
            )
            use_ai = gr.Checkbox(
                label="Use HF AI enhancement when HF_TOKEN is configured",
                value=True,
            )
            generate_btn = gr.Button("✨ Generate Prompt", variant="primary")

        with gr.Column(scale=2):
            output = gr.Textbox(
                label="Generated prompt",
                lines=20,
                show_copy_button=True,
            )

    generate_btn.click(
        generate,
        inputs=[idea, style, aspect, mood, detail, camera, use_ai],
        outputs=output,
    )

    gr.Examples(
        examples=[
            ["A futuristic Brisbane skyline after a tropical storm", "Cinematic", "16:9", "moody and atmospheric", "highly detailed", "wide establishing shot"],
            ["A tiny robot repairing an old mechanical clock", "3D Render", "1:1", "warm and whimsical", "ultra detailed", "close-up macro composition"],
            ["A medieval library hidden inside a giant tree", "Fantasy", "4:5", "magical and serene", "richly detailed", "symmetrical interior composition"],
        ],
        inputs=[idea, style, aspect, mood, detail, camera],
    )

demo.launch()