# /app.py
"""
Midjourney Prompt Expander —V 7

NSFW filter (strict/relaxed), Negative Prompt (--no), MJ v7 AR presets.
Quality presets: --stylize, --quality, --chaos (toggles + values).
History panel: auto-save results, table, CSV export, clear.
Radio presets: Product / Portrait / Scene -> auto-sets AR + quality flags.
Copy-ready /imagine command. Optional input -> random concept.
"""

from __future__ import annotations

import os
import re
import uuid
import random
import textwrap
from datetime import datetime
from typing import List, Dict, Tuple, Set, Any

import gradio as gr
import pandas as pd

# ---------- Prompt parts ----------
SUBJECTS = [
    "ancient tree with glowing runes", "cyberpunk alley at midnight", "desert caravan crossing dunes",
    "bioluminescent jellyfish in the deep sea", "abandoned spacecraft interior", "steampunk airship above clouds",
    "snowy mountain temple", "mystic librarian surrounded by floating books", "futuristic samurai in neon rain",
    "retro diner at sunrise", "miniature city inside a glass terrarium", "astronaut tending a cosmic garden",
]
ACTIONS = [
    "emerging from fog", "caught mid-motion", "levitating gently", "cracking through the surface",
    "reflecting in a puddle", "silhouetted by the sun", "bathed in moonlight", "draped in shadows",
    "exploding with color", "frozen in time",
]
ENVIRONMENTS = [
    "narrow cobblestone streets", "vast open steppe", "misty pine forest", "glittering coastline",
    "industrial refinery", "crystalline cavern", "stormy sky", "quiet museum gallery",
]
STYLES = [
    "photorealistic", "cinematic", "studio lighting", "documentary", "surreal", "dreamlike",
    "isometric", "illustrative", "watercolor", "oil painting", "low-poly", "line art",
]
LIGHTING = [
    "golden hour", "soft diffused light", "backlit", "rim lighting", "hard contrast", "overcast",
    "volumetric light rays", "neon glow", "bioluminescent glow", "subsurface scattering",
]
COMPOSITION = [
    "rule of thirds", "center framing", "leading lines", "wide shot", "close-up macro",
    "aerial view", "dutch angle", "symmetrical balance", "shallow depth of field",
]
COLOR_PALETTES = [
    "teal and orange", "muted pastels", "high-contrast monochrome", "warm sepia",
    "vibrant neon", "earthy tones", "cool blues and violets",
]
LENSES = ["35mm", "50mm", "85mm", "24mm", "macro 100mm", "anamorphic", "tilt-shift"]
QUALITY_TAGS = [
    "ultra-detailed", "high dynamic range", "8k", "physically based rendering", "ray-traced reflections",
    "textured realism", "film grain", "bokeh",
]
ART_REFERENCES = [
    "artstation trending", "award-winning", "magazine editorial style", "studio ghibli vibes",
    "moebius-inspired linework", "greg rutkowski style", "ian hubert lighting",
]

# ---------- MJ v7 aspect-ratio presets ----------
ASPECT_RATIOS_MJ7: Dict[str, str] = {
    "1:1 (square)": "--ar 1:1",
    "3:2 (landscape)": "--ar 3:2",
    "2:3 (portrait)": "--ar 2:3",
    "16:9 (widescreen)": "--ar 16:9",
    "9:16 (vertical)": "--ar 9:16",
    "4:5 (portrait)": "--ar 4:5",
    "5:4 (landscape)": "--ar 5:4",
    "21:9 (ultrawide)": "--ar 21:9",
}

MJ_VERSION_SUFFIX = " —V 7"  # exact suffix with an em dash

# ---------- NSFW filtering ----------
# Minimal deterministic lists. Avoid real slurs in code.
SEXUAL: Set[str] = {
    "nude", "naked", "nsfw", "porn", "erotic", "sexual", "sensual", "fetish", "kink", "explicit",
    "strip", "lingerie", "breasts", "nipples", "genitals", "penis", "vagina", "hentai", "bdsm",
}
MINORS: Set[str] = {"child", "children", "kid", "minor", "teen", "young girl", "young boy", "loli", "shota"}
VIOLENCE: Set[str] = {
    "gore", "gory", "bloodshed", "decapitated", "severed", "dismembered", "eviscerated", "torture",
}
SLURS: Set[str] = {"slur_placeholder_1", "slur_placeholder_2"}  # placeholders only

SANITIZE_MAP: Dict[str, str] = {
    "nude": "tasteful minimal attire",
    "naked": "minimal attire",
    "lingerie": "fashion attire",
    "breasts": "torso",
    "nipples": "torso",
    "genitals": "body",
    "penis": "body",
    "vagina": "body",
    "porn": "artistic",
    "erotic": "romantic",
    "fetish": "stylized",
    "kink": "stylized",
    "explicit": "artistic",
    "gore": "dramatic",
    "gory": "dramatic",
    "bloodshed": "intense",
    "decapitated": "abstract",
    "dismembered": "abstract",
    "eviscerated": "abstract",
    "torture": "dark",
}

WORD_RE = re.compile(r"[a-z0-9]+(?:'[a-z0-9]+)?", re.IGNORECASE)

DEFAULT_NEGATIVE_LIST = (
    "blurry, low-res, low quality, jpeg artifacts, watermark, signature, text, logo, "
    "overexposed, underexposed, noisy, out of frame, cropped, duplicate, "
    "deformed, disfigured, mutated, extra limbs, extra fingers, poorly drawn, bad anatomy"
)

# ---------- Radio presets ----------
PRESETS: Dict[str, Dict[str, Any]] = {
    "Product": {
        "aspect": "4:5 (portrait)",
        "stylize_enabled": True, "stylize_value": 100,
        "quality_enabled": True, "quality_value": 2.0,
        "chaos_enabled": False, "chaos_value": 0,
    },
    "Portrait": {
        "aspect": "2:3 (portrait)",
        "stylize_enabled": True, "stylize_value": 100,
        "quality_enabled": True, "quality_value": 1.0,
        "chaos_enabled": False, "chaos_value": 0,
    },
    "Scene": {
        "aspect": "16:9 (widescreen)",
        "stylize_enabled": True, "stylize_value": 250,
        "quality_enabled": True, "quality_value": 1.0,
        "chaos_enabled": True, "chaos_value": 10,
    },
}

# ---------- Helpers ----------
def _tokenize(text: str) -> List[str]:
    return [m.group(0).lower() for m in WORD_RE.finditer(text or "")]

def detect_nsfw(text: str) -> Dict[str, List[str]]:
    toks = set(_tokenize(text))
    hits: Dict[str, List[str]] = {}
    cats = {"sexual": SEXUAL, "minors": MINORS, "violence": VIOLENCE, "slurs": SLURS}
    for name, vocab in cats.items():
        matched = sorted(toks.intersection(vocab))
        if matched:
            hits[name] = matched
    return hits

def sanitize_text(text: str) -> str:
    if not text:
        return text

    def repl(match: re.Match) -> str:
        w = match.group(0)
        lw = w.lower()
        if lw in SANITIZE_MAP:
            rep = SANITIZE_MAP[lw]
            # Preserve capitalization of first letter for readability
            return rep.capitalize() if w and w[0].isupper() else rep
        return w

    pattern = re.compile(r"\b(" + "|".join(map(re.escape, SANITIZE_MAP.keys())) + r")\b", re.IGNORECASE)
    return pattern.sub(repl, text)

def _rng(seed: int | None) -> random.Random:
    return random.Random(seed) if seed is not None else random.Random()

def _pick(r: random.Random, items: List[str], k: int = 1) -> List[str]:
    if k <= 0:
        return []
    if k >= len(items):
        items = items[:]
        r.shuffle(items)
        return items
    return r.sample(items, k)

def synthesize_subject(r: random.Random) -> str:
    base = r.choice(SUBJECTS)
    maybe = f"{base}, {r.choice(ACTIONS)}" if r.random() < 0.85 else base
    if r.random() < 0.7:
        maybe += f", {r.choice(ENVIRONMENTS)}"
    return maybe

def expand_prompt(
    idea: str | None,
    style_hint: str,
    aspect_ratio_flag: str,
    creativity: float,
    seed: int | None,
    negative_text: str | None,
    stylize_enabled: bool,
    stylize_value: int,
    quality_enabled: bool,
    quality_value: float,
    chaos_enabled: bool,
    chaos_value: int,
) -> str:
    """Build expanded MJ prompt, append --no (if any), presets, and —V 7."""
    r = _rng(seed)
    base = (idea or "").strip()
    subject = base if base else synthesize_subject(r)

    k_style = 1 + int(creativity * 1.0)
    k_light = 1 + int(creativity * 1.0)
    k_comp  = 1 + int(creativity * 1.0)
    k_color = 1 if creativity < 0.66 else 2
    k_quality = 2 + int(creativity * 2.0)

    styles = [style_hint] if style_hint != "(auto)" else _pick(r, STYLES, k_style)
    lights = _pick(r, LIGHTING, k_light)
    comps  = _pick(r, COMPOSITION, k_comp)
    colors = _pick(r, COLOR_PALETTES, k_color)
    lenses = _pick(r, LENSES, 1 if r.random() < 0.8 else 2)
    quality = _pick(r, QUALITY_TAGS, min(k_quality, len(QUALITY_TAGS)))
    refs   = _pick(r, ART_REFERENCES, 1 if r.random() < 0.6 else 2)

    parts = [
        subject,
        ", ".join(styles),
        ", ".join(lights),
        ", ".join(comps),
        ", ".join(colors),
        f"{', '.join(lenses)} lens",
        ", ".join(refs),
        ", ".join(quality),
        aspect_ratio_flag,
    ]

    neg = (negative_text or "").strip()
    if neg:
        parts.append(f"--no {neg}")

    if stylize_enabled:
        parts.append(f"--stylize {int(stylize_value)}")
    if quality_enabled:
        q_str = str(quality_value).rstrip("0").rstrip(".")
        parts.append(f"--quality {q_str}")
    if chaos_enabled:
        parts.append(f"--chaos {int(chaos_value)}")

    prompt = ", ".join([p for p in parts if p and p.strip()])
    prompt = textwrap.fill(prompt + MJ_VERSION_SUFFIX, width=120)
    return prompt

# ---------- History ----------
HISTORY_COLUMNS = [
    "time", "preset", "idea", "style", "aspect", "creativity", "seed",
    "negatives", "flags", "expanded", "command"
]

def history_to_df(history: List[Dict[str, Any]]) -> pd.DataFrame:
    if not history:
        return pd.DataFrame(columns=HISTORY_COLUMNS)
    return pd.DataFrame(history, columns=HISTORY_COLUMNS)

def export_history_csv(history: List[Dict[str, Any]]) -> str:
    """Create a temp CSV and return its path."""
    df = history_to_df(history)
    tmp_dir = "/tmp"
    os.makedirs(tmp_dir, exist_ok=True)
    path = os.path.join(tmp_dir, f"prompts_{uuid.uuid4().hex[:8]}.csv")
    df.to_csv(path, index=False)
    return path

# ---------- Gradio bridges ----------
def generate_with_safety(
    idea: str,
    style: str,
    preset_name: str,
    aspect: str,
    creativity: float,
    seed_text: str,
    strict_mode: bool,
    negative_text: str,
    use_default_negative_when_empty: bool,
    stylize_enabled: bool,
    stylize_value: int,
    quality_enabled: bool,
    quality_value: float,
    chaos_enabled: bool,
    chaos_value: int,
    auto_save_history: bool,
    history: List[Dict[str, Any]],
    max_retries: int = 3,
) -> Tuple[str, str, List[Dict[str, Any]], pd.DataFrame]:
    """Generate expanded prompt + MJ command with safety, presets, and history."""
    # Seed
    seed: int | None = None
    if seed_text and seed_text.strip():
        try:
            seed = int(seed_text.strip())
        except ValueError:
            seed = None  # ignore invalid seed

    # Negatives
    neg_src = (negative_text or "").strip()
    if not neg_src and use_default_negative_when_empty:
        neg_src = DEFAULT_NEGATIVE_LIST

    # Input safety
    if strict_mode:
        if idea:
            hits = detect_nsfw(idea)
            if hits:
                cats = ", ".join(f"{k}: {', '.join(v)}" for k, v in hits.items())
                return (f"Blocked by NSFW filter (input contains {cats}). Try different wording.", "", history, history_to_df(history))
        if neg_src:
            hits_neg = detect_nsfw(neg_src)
            if hits_neg:
                cats = ", ".join(f"{k}: {', '.join(v)}" for k, v in hits_neg.items())
                return (f"Blocked by NSFW filter (negative prompt contains {cats}). Edit negatives.", "", history, history_to_df(history))
        safe_idea, safe_neg = idea, neg_src
    else:
        safe_idea = sanitize_text(idea) if idea else idea
        safe_neg = sanitize_text(neg_src) if neg_src else neg_src

    # Generate (retry if blocked in strict)
    rseed = seed
    for _ in range(max_retries):
        expanded = expand_prompt(
            safe_idea,
            style,
            ASPECT_RATIOS_MJ7[aspect],
            creativity,
            rseed,
            safe_neg,
            stylize_enabled,
            stylize_value,
            quality_enabled,
            quality_value,
            chaos_enabled,
            chaos_value,
        )
        hits_out = detect_nsfw(expanded)
        if not hits_out:
            mj_cmd = f"/imagine prompt: {expanded}"

            # Auto-save to history
            if auto_save_history:
                flags_str = ", ".join(
                    x for x in [
                        f"--stylize {stylize_value}" if stylize_enabled else None,
                        f"--quality {str(quality_value).rstrip('0').rstrip('.')}" if quality_enabled else None,
                        f"--chaos {chaos_value}" if chaos_enabled else None,
                        f"--no {safe_neg}" if safe_neg else None,
                        ASPECT_RATIOS_MJ7[aspect],
                        "—V 7",
                    ] if x
                )
                history.append({
                    "time": datetime.utcnow().isoformat(timespec="seconds") + "Z",
                    "preset": preset_name,
                    "idea": idea,
                    "style": style,
                    "aspect": aspect,
                    "creativity": creativity,
                    "seed": seed if seed is not None else "",
                    "negatives": safe_neg,
                    "flags": flags_str,
                    "expanded": expanded,
                    "command": mj_cmd,
                })

            return expanded, mj_cmd, history, history_to_df(history)

        if strict_mode:
            rseed = random.randint(0, 2**31 - 1)
            continue
        else:
            expanded_sanitized = sanitize_text(expanded)
            if not detect_nsfw(expanded_sanitized):
                mj_cmd = f"/imagine prompt: {expanded_sanitized}"
                if auto_save_history:
                    history.append({
                        "time": datetime.utcnow().isoformat(timespec="seconds") + "Z",
                        "preset": preset_name,
                        "idea": idea,
                        "style": style,
                        "aspect": aspect,
                        "creativity": creativity,
                        "seed": seed if seed is not None else "",
                        "negatives": safe_neg,
                        "flags": "(sanitized)",
                        "expanded": expanded_sanitized,
                        "command": mj_cmd,
                    })
                return expanded_sanitized, mj_cmd, history, history_to_df(history)
            rseed = random.randint(0, 2**31 - 1)

    return ("Blocked by NSFW filter after multiple attempts. Please adjust your idea or disable strict mode.",
            "", history, history_to_df(history))

def apply_preset(preset_name: str):
    """Return component values for AR + quality flags based on radio preset."""
    p = PRESETS[preset_name]
    return (
        p["aspect"],
        gr.update(value=p["stylize_enabled"]), p["stylize_value"],
        gr.update(value=p["quality_enabled"]), p["quality_value"],
        gr.update(value=p["chaos_enabled"]), p["chaos_value"],
    )

def clear_history(history: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], pd.DataFrame, gr.File]:
    """Clear in-memory history and hide export file."""
    history.clear()
    return history, history_to_df(history), gr.update(visible=False, value=None)

def do_export(history: List[Dict[str, Any]]):
    """Create CSV file and expose it in the UI."""
    if not history:
        return gr.update(value=None, visible=False)
    path = export_history_csv(history)
    return gr.update(value=path, visible=True)

# ---------- UI ----------
with gr.Blocks(title="Midjourney Prompt Expander —V 7", fill_height=True) as demo:
    gr.Markdown("# ✨ Midjourney Prompt Expander\nShort idea in → rich prompt out. (Always appends —V 7.)")

    history_state = gr.State([])  # list[dict]

    with gr.Row():
        idea = gr.Textbox(
            label="Your idea (optional)",
            placeholder="e.g., 'a cozy reading nook' or leave empty for a random concept",
            lines=2,
        )

    with gr.Row():
        style = gr.Dropdown(
            label="Style",
            choices=["(auto)"] + STYLES,
            value="(auto)",
            allow_custom_value=True,
            info="Pick a style or type your own.",
        )

        preset_radio = gr.Radio(
            label="Preset",
            choices=list(PRESETS.keys()),
            value="Scene",
            info="Auto-sets AR + quality flags.",
        )

        aspect = gr.Dropdown(
            label="Aspect Ratio (MJ v7 presets)",
            choices=list(ASPECT_RATIOS_MJ7.keys()),
            value="16:9 (widescreen)",
        )

    with gr.Row():
        creativity = gr.Slider(
            label="Creativity",
            minimum=0.0, maximum=1.0, value=0.6, step=0.05,
            info="Higher = more modifiers.",
        )
        seed = gr.Textbox(
            label="Seed (optional, integer)",
            placeholder="e.g., 12345 for reproducibility",
        )

    with gr.Row():
        negative_text = gr.Textbox(
            label="Negative prompt (--no ...)",
            placeholder="e.g., blurry, watermark, text, low-res",
            lines=2,
        )

    with gr.Row():
        use_default_negative = gr.Checkbox(
            label="Auto-insert default negative list when empty",
            value=True,
        )
        strict_mode = gr.Checkbox(
            label="NSFW Strict Mode",
            value=True,
            info="Strict: block unsafe words. Off: sanitize instead.",
        )

    with gr.Accordion("Quality Flags", open=True):
        with gr.Row():
            stylize_enabled = gr.Checkbox(label="--stylize", value=True)
            stylize_value = gr.Slider(label="stylize value", minimum=0, maximum=1000, step=10, value=250)
        with gr.Row():
            quality_enabled = gr.Checkbox(label="--quality", value=True)
            quality_value = gr.Slider(label="quality value", minimum=0.25, maximum=2.0, step=0.25, value=1.0)
        with gr.Row():
            chaos_enabled = gr.Checkbox(label="--chaos", value=False)
            chaos_value = gr.Slider(label="chaos value", minimum=0, maximum=100, step=1, value=10)
        auto_save_history = gr.Checkbox(label="Auto-save results to history", value=True)

    with gr.Row():
        generate_btn = gr.Button("Generate", variant="primary")
        clear_btn = gr.Button("Clear History")
        export_btn = gr.Button("Export CSV")
        export_file = gr.File(label="Download CSV", visible=False)

    with gr.Row():
        expanded = gr.Textbox(
            label="Expanded Prompt (appends —V 7)",
            lines=6,
            show_copy_button=True,
        )
    with gr.Row():
        mj_command = gr.Textbox(
            label="Copy-ready /imagine command",
            value="",
            lines=4,
            show_copy_button=True,
        )

    gr.Markdown("### History")
    history_df = gr.Dataframe(
        value=pd.DataFrame(columns=HISTORY_COLUMNS),
        column_names=HISTORY_COLUMNS,
        interactive=False,
        height=300,
    )

    # --- Wiring ---
    preset_radio.change(
        fn=apply_preset,
        inputs=[preset_radio],
        outputs=[aspect, stylize_enabled, stylize_value, quality_enabled, quality_value, chaos_enabled, chaos_value],
    )

    generate_btn.click(
        fn=generate_with_safety,
        inputs=[
            idea,              # str
            style,             # str
            preset_radio,      # str (name)
            aspect,            # str (key to AR map)
            creativity,        # float
            seed,              # str (maybe empty)
            strict_mode,       # bool
            negative_text,     # str
            use_default_negative,  # bool
            stylize_enabled,   # bool
            stylize_value,     # int
            quality_enabled,   # bool
            quality_value,     # float
            chaos_enabled,     # bool
            chaos_value,       # int
            auto_save_history, # bool
            history_state,     # list[dict]
        ],
        outputs=[
            expanded,          # expanded prompt
            mj_command,        # copy-ready command
            history_state,     # updated history state
            history_df,        # updated dataframe
        ],
    )

    clear_btn.click(
        fn=clear_history,
        inputs=[history_state],
        outputs=[history_state, history_df, export_file],
    )

    export_btn.click(
        fn=do_export,
        inputs=[history_state],
        outputs=[export_file],
    )

# For local dev; on HF Spaces, the runner calls `demo` automatically.
if __name__ == "__main__":
    demo.launch()