#!/usr/bin/env python3
"""
Marketing poster/image generation via SophNet Qwen-Image task API.

Generates marketing posters and promotional images for various platforms
(WeChat Moments, Xiaohongshu, WeChat Official Account, etc.) using async
task polling. Supports style presets, fine-grained style dimensions,
text overlay, and reference-image paths (prompt guidance only).

Usage:
    uv run {baseDir}/scripts/generate_poster.py --type moments --prompt "..." [--style-preset promo]
    uv run {baseDir}/scripts/generate_poster.py --type xiaohongshu --style-preset kawaii --prompt "..."
    uv run {baseDir}/scripts/generate_poster.py --type moments --prompt "..." --text-overlay '{"title":"母亲节特惠"}'
    uv run {baseDir}/scripts/generate_poster.py --type moments --prompt "..." --reference-image product.jpg
"""

from __future__ import annotations

import argparse
import json
import math
import sys
import time
from typing import Any, Dict, List, Optional, Tuple

import requests

# ---------------------------------------------------------------------------
# API
# ---------------------------------------------------------------------------

API_URL = "https://www.sophnet.com/api/open-apis/projects/easyllms/imagegenerator/task"
DEFAULT_MODEL = "Qwen-Image"

# Qwen-Image supported sizes (pick nearest aspect ratio)
QWEN_SIZES: List[Tuple[int, int]] = [
    (1328, 1328),   # 1:1
    (1140, 1472),   # ~3:4
    (1472, 1140),   # ~4:3
    (1664, 928),    # ~16:9
    (928, 1664),    # ~9:16
    (1024, 1024),
]

# ---------------------------------------------------------------------------
# Safety
# ---------------------------------------------------------------------------

SAFETY_NEGATIVE_PROMPT = (
    "nsfw, nudity, nude, naked, sexual, erotic, pornographic, gore, blood, violence, "
    "bloody, corpse, dead body, weapon, gun, knife, drugs, smoking, alcohol, gambling, "
    "politically sensitive, national flag, national emblem, political leader, "
    "religious symbol, hate symbol, discrimination, racist, offensive, disturbing, "
    "child exploitation, terrorism, self-harm"
)

# ---------------------------------------------------------------------------
# Poster types
# ---------------------------------------------------------------------------

POSTER_TYPES: Dict[str, Dict[str, str]] = {
    "moments": {
        "size": "1080*1080",
        "aspect_ratio": "1:1",
        "label": "朋友圈海报 (1080×1080)",
    },
    "xiaohongshu": {
        "size": "1080*1440",
        "aspect_ratio": "3:4",
        "label": "小红书封面 (1080×1440)",
    },
    "wechat-header": {
        "size": "900*383",
        "aspect_ratio": "16:9",
        "label": "公众号头图 (900×383)",
    },
    "wechat-square": {
        "size": "200*200",
        "aspect_ratio": "1:1",
        "label": "公众号方形预览 (200×200)",
    },
    "share-card": {
        "size": "500*400",
        "aspect_ratio": "4:3",
        "label": "微信群分享图 (500×400)",
    },
    "product": {
        "size": "1024*1024",
        "aspect_ratio": "1:1",
        "label": "产品展示图 (1024×1024)",
    },
    "guide": {
        "size": "1080*1440",
        "aspect_ratio": "3:4",
        "label": "攻略图/信息图 (1080×1440)",
    },
}

# ---------------------------------------------------------------------------
# Style dimensions
# ---------------------------------------------------------------------------

PALETTES: Dict[str, str] = {
    "warm": "Warm color palette: golden yellows, amber oranges, terracotta reds, honey tones. Evokes comfort and warmth.",
    "elegant": "Elegant palette: champagne gold, ivory, dusty rose, soft grey. Refined and sophisticated feel.",
    "cool": "Cool palette: ocean blues, mint greens, silver greys, icy whites. Clean and calming atmosphere.",
    "dark": "Dark palette: deep navy, charcoal, dark teal, muted burgundy. High contrast, dramatic mood.",
    "earth": "Earth palette: olive green, clay brown, sandstone, forest tones. Natural and organic feeling.",
    "vivid": "Vivid palette: saturated primary colors, bold contrasts, bright accents. High energy and eye-catching.",
    "pastel": "Pastel palette: soft pink, baby blue, lavender, mint. Gentle, dreamy, and delicate.",
    "mono": "Monochrome palette: shades of a single hue with tonal variation. Unified and striking.",
    "retro": "Retro palette: mustard yellow, burnt orange, avocado green, faded teal. 1970s nostalgic warmth.",
}

RENDERINGS: Dict[str, str] = {
    "flat-vector": "Flat vector illustration style: clean geometric shapes, solid color fills, no gradients or textures, minimal line work.",
    "hand-drawn": "Hand-drawn illustration style: visible sketch lines, organic imperfections, ink-and-paper feel, slightly uneven edges.",
    "painterly": "Painterly style: visible brush strokes, rich textures, soft blended edges, oil or watercolor painting feel.",
    "digital": "Polished digital art style: smooth gradients, clean rendering, precise details, modern and refined.",
    "pixel": "Pixel art style: retro 8-bit aesthetic, blocky shapes, limited color palette, nostalgic gaming feel.",
    "chalk": "Chalk/chalkboard style: white and colored chalk on dark background, hand-lettered feel, educational aesthetic.",
}

MOODS: Dict[str, str] = {
    "subtle": "Subtle mood: low contrast, muted tones, generous whitespace, understated and calm composition.",
    "balanced": "Balanced mood: moderate contrast, harmonious composition, professional and approachable.",
    "bold": "Bold mood: high contrast, saturated colors, dynamic composition, strong visual impact.",
}

LAYOUTS: Dict[str, str] = {
    "bento-grid": "Bento grid layout: modular grid of varied-size cards, each containing a distinct piece of information.",
    "list": "List layout: enumerated items in a vertical sequence with icons or numbers.",
    "comparison": "Comparison layout: side-by-side or split-screen contrasting two options or concepts.",
    "flow": "Flow layout: connected steps or stages showing a process or timeline with directional arrows.",
    "mindmap": "Mind map layout: central concept radiating outward to connected branches and sub-topics.",
    "hub-spoke": "Hub-spoke layout: central element surrounded by related items in a radial arrangement.",
    "funnel": "Funnel layout: wide-to-narrow stages showing progressive filtering or conversion.",
    "dense-modules": "Dense modules layout: tightly packed information blocks with high data density, guide-style.",
}

STYLE_PRESETS: Dict[str, Dict[str, str]] = {
    "promo": {"palette": "vivid", "rendering": "flat-vector", "mood": "bold"},
    "elegant": {"palette": "earth", "rendering": "hand-drawn", "mood": "subtle"},
    "minimal": {"palette": "mono", "rendering": "flat-vector", "mood": "subtle"},
    "festive": {"palette": "warm", "rendering": "painterly", "mood": "bold"},
    "cozy": {"palette": "warm", "rendering": "painterly", "mood": "subtle"},
    "kawaii": {"palette": "pastel", "rendering": "flat-vector", "mood": "balanced"},
    "morandi": {"palette": "earth", "rendering": "hand-drawn", "mood": "subtle"},
    "pop-art": {"palette": "vivid", "rendering": "flat-vector", "mood": "bold"},
    "vintage": {"palette": "retro", "rendering": "hand-drawn", "mood": "balanced"},
    "blueprint": {"palette": "dark", "rendering": "chalk", "mood": "bold"},
    "notion": {"palette": "mono", "rendering": "hand-drawn", "mood": "subtle"},
    "watercolor": {"palette": "pastel", "rendering": "painterly", "mood": "subtle"},
    "corporate": {"palette": "cool", "rendering": "flat-vector", "mood": "balanced"},
}

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def parse_size(size_str: str) -> Tuple[int, int]:
    """Parse '1080x1440', '1080*1440', or '1080×1440' into (width, height)."""
    normalized = size_str.replace("x", "*").replace("×", "*")
    parts = normalized.split("*")
    if len(parts) != 2:
        raise ValueError(f"Invalid size format: {size_str!r}. Use WxH (e.g. 1080x1440).")
    return int(parts[0].strip()), int(parts[1].strip())


def size_to_qwen_size(w: int, h: int) -> str:
    """Map requested dimensions to nearest Qwen-Image supported size."""
    target = w / h
    best_w, best_h = QWEN_SIZES[0]
    best_diff = float("inf")
    for qw, qh in QWEN_SIZES:
        diff = abs(math.log(target) - math.log(qw / qh))
        if diff < best_diff:
            best_diff = diff
            best_w, best_h = qw, qh
    return f"{best_w}*{best_h}"


def resolve_style(
    preset: Optional[str],
    palette: Optional[str],
    rendering: Optional[str],
    mood: Optional[str],
) -> Tuple[Optional[str], Optional[str], Optional[str]]:
    """Apply preset defaults, then override with explicit values."""
    p_pal, p_ren, p_mood = None, None, None
    if preset and preset in STYLE_PRESETS:
        cfg = STYLE_PRESETS[preset]
        p_pal, p_ren, p_mood = cfg.get("palette"), cfg.get("rendering"), cfg.get("mood")
    return (palette or p_pal, rendering or p_ren, mood or p_mood)


MAX_REFERENCE_IMAGES = 10


def validate_reference_image_paths(paths: List[str]) -> List[str]:
    """Validate local reference image paths exist."""
    from pathlib import Path

    resolved: List[str] = []
    for path in paths:
        p = Path(path).expanduser()
        if not p.exists():
            raise FileNotFoundError(f"Reference image not found: {path}")
        resolved.append(str(p))
    return resolved


# ---------------------------------------------------------------------------
# Prompt construction
# ---------------------------------------------------------------------------


def build_prompt(
    user_prompt: str,
    width: int,
    height: int,
    palette: Optional[str] = None,
    rendering: Optional[str] = None,
    mood: Optional[str] = None,
    layout: Optional[str] = None,
    text_overlay: Optional[Dict[str, str]] = None,
    reference_image_paths: Optional[List[str]] = None,
) -> Tuple[str, str]:
    """Return (prompt, negative_prompt) for Qwen-Image API."""
    merged_negative = SAFETY_NEGATIVE_PROMPT

    parts = [
        "Generate a marketing poster image with these specifications:",
        f"- Desired dimensions: {width}x{height} pixels",
    ]

    if reference_image_paths:
        names = ", ".join(reference_image_paths)
        parts.append(
            "- The user provided reference materials (product photos, logos, etc.) at: "
            f"{names}. Faithfully depict the products, logos, or brand elements described "
            "in the scene prompt with accurate colors, shapes, and visual details."
        )

    if palette and palette in PALETTES:
        parts.append(f"- Color direction: {PALETTES[palette]}")
    if rendering and rendering in RENDERINGS:
        parts.append(f"- Rendering style: {RENDERINGS[rendering]}")
    if mood and mood in MOODS:
        parts.append(f"- Mood: {MOODS[mood]}")
    if layout and layout in LAYOUTS:
        parts.append(f"- Information layout: {LAYOUTS[layout]}")

    if text_overlay:
        text_parts = []
        if text_overlay.get("title"):
            text_parts.append(f'main title text "{text_overlay["title"]}"')
        if text_overlay.get("subtitle"):
            text_parts.append(f'subtitle text "{text_overlay["subtitle"]}"')
        if text_overlay.get("price"):
            text_parts.append(f'prominent price tag "{text_overlay["price"]}"')
        if text_overlay.get("footer"):
            text_parts.append(f'footer text "{text_overlay["footer"]}"')
        if text_parts:
            parts.append(
                f"- Text to render on the image: {', '.join(text_parts)}. "
                "Text must be clearly legible, well-positioned, and not overlap key visual elements."
            )

    parts.append(f"- {user_prompt}")
    return "\n".join(parts), merged_negative


# ---------------------------------------------------------------------------
# API interaction
# ---------------------------------------------------------------------------


def _extract_task_id(data: Dict[str, Any]) -> Optional[str]:
    for key in ("task_id", "taskId", "taskID", "id"):
        if key in data:
            return str(data[key])
    output = data.get("output")
    if isinstance(output, dict):
        for key in ("task_id", "taskId", "taskID", "id"):
            if key in output:
                return str(output[key])
    return None


def create_task(
    api_key: str,
    prompt: str,
    size: str,
    negative_prompt: Optional[str] = None,
    model: str = DEFAULT_MODEL,
) -> str:
    payload: Dict[str, Any] = {
        "model": model,
        "input": {"prompt": prompt},
        "parameters": {
            "size": size,
            "n": 1,
            "watermark": False,
            "prompt_extend": True,
        },
    }
    if negative_prompt:
        payload["input"]["negative_prompt"] = negative_prompt

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    resp = requests.post(API_URL, json=payload, headers=headers, timeout=60)
    resp.raise_for_status()
    data = resp.json()

    task_id = _extract_task_id(data)
    if not task_id:
        raise ValueError(f"task_id not found in response: {json.dumps(data, ensure_ascii=False)}")
    return task_id


def poll_task(
    api_key: str,
    task_id: str,
    poll_interval: int = 2,
    max_wait: int = 300,
) -> Dict[str, Any]:
    headers = {"Authorization": f"Bearer {api_key}"}
    start = time.time()

    while True:
        if time.time() - start > max_wait:
            raise TimeoutError(f"timed out after {max_wait}s waiting for task {task_id}")

        resp = requests.get(f"{API_URL}/{task_id}", headers=headers, timeout=60)
        resp.raise_for_status()
        data = resp.json()

        status = None
        for key in ("status", "taskStatus", "task_status"):
            if key in data:
                status = data[key]
                break
            output = data.get("output")
            if isinstance(output, dict) and key in output:
                status = output[key]
                break

        if status and status.lower() in ("succeeded", "success"):
            return data
        if status and status.lower() in ("failed", "error"):
            raise RuntimeError(json.dumps(data, ensure_ascii=False))

        time.sleep(poll_interval)


def extract_image_urls(data: Dict[str, Any]) -> List[str]:
    urls: List[str] = []

    def _walk(obj: Any) -> None:
        if isinstance(obj, dict):
            if "url" in obj and isinstance(obj["url"], str):
                urls.append(obj["url"])
            for v in obj.values():
                _walk(v)
        elif isinstance(obj, list):
            for item in obj:
                _walk(item)

    _walk(data)
    return urls


def generate_image(
    api_key: str,
    prompt: str,
    size: str,
    negative_prompt: Optional[str] = None,
    model: str = DEFAULT_MODEL,
    poll_interval: int = 2,
    max_wait: int = 300,
) -> Tuple[str, str]:
    """Create task, poll until done, return (task_id, image_url)."""
    task_id = create_task(api_key, prompt, size, negative_prompt=negative_prompt, model=model)
    result = poll_task(api_key, task_id, poll_interval=poll_interval, max_wait=max_wait)
    urls = extract_image_urls(result)
    if not urls:
        raise ValueError(f"url not found in response: {json.dumps(result, ensure_ascii=False)}")
    return task_id, urls[0]


def handle_http_error(e: requests.HTTPError) -> None:
    status_code = e.response.status_code if e.response is not None else None
    if status_code in (401, 403):
        print("STATUS=permission_denied", file=sys.stderr)
        print(
            "Error: 当前账户没有所调用图片生成模型的使用权限。"
            "请联系 Sophclaw 平台客服开通 Qwen-Image 权限后重试。",
            file=sys.stderr,
        )
    else:
        print("STATUS=api_error", file=sys.stderr)
        print(
            f"Error: SophNet Qwen-Image API call failed (HTTP {status_code}): {e}",
            file=sys.stderr,
        )
    sys.exit(1)


# ---------------------------------------------------------------------------
# Output formatting
# ---------------------------------------------------------------------------


def format_output(
    poster_type: str,
    w: int,
    h: int,
    api_size: str,
    palette: Optional[str],
    rendering: Optional[str],
    mood: Optional[str],
    layout: Optional[str],
    style_preset: Optional[str],
    status: str,
    image_url: Optional[str],
    task_id: Optional[str] = None,
    reference_image_count: int = 0,
    fmt: str = "text",
) -> str:
    data: Dict[str, str] = {
        "POSTER_TYPE": poster_type,
        "POSTER_SIZE": f"{w}*{h}",
        "API_SIZE": api_size,
        "MODEL": DEFAULT_MODEL,
    }
    if style_preset:
        data["STYLE_PRESET"] = style_preset
    if palette:
        data["PALETTE"] = palette
    if rendering:
        data["RENDERING"] = rendering
    if mood:
        data["MOOD"] = mood
    if layout:
        data["LAYOUT"] = layout
    if reference_image_count:
        data["REFERENCE_IMAGES"] = str(reference_image_count)
    if task_id:
        data["TASK_ID"] = task_id
    data["STATUS"] = status
    if image_url:
        data["IMAGE_URL"] = image_url

    if fmt == "json":
        return json.dumps(data, ensure_ascii=False, indent=2)
    return "\n".join(f"{k}={v}" for k, v in data.items())


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main(argv: Optional[List[str]] = None) -> None:
    type_names = ", ".join(f"{k} ({v['label']})" for k, v in POSTER_TYPES.items())
    preset_names = ", ".join(STYLE_PRESETS.keys())

    parser = argparse.ArgumentParser(
        description="Generate marketing posters via SophNet Qwen-Image API.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            f"Poster types:\n  {type_names}\n\n"
            f"Style presets (shorthand for palette+rendering+mood):\n  {preset_names}\n\n"
            f"Palettes:\n  {', '.join(PALETTES.keys())}\n\n"
            f"Renderings:\n  {', '.join(RENDERINGS.keys())}\n\n"
            f"Moods:\n  {', '.join(MOODS.keys())}\n\n"
            f"Layouts (guide type only):\n  {', '.join(LAYOUTS.keys())}"
        ),
    )
    parser.add_argument("--type", required=True, choices=POSTER_TYPES.keys(),
                        help="Poster type (determines default size and aspect ratio)")
    parser.add_argument("--prompt", required=True,
                        help="Image content description (scene, subject, atmosphere)")
    parser.add_argument("--style-preset", default=None, choices=list(STYLE_PRESETS.keys()),
                        help="Style shorthand that sets palette+rendering+mood together")
    parser.add_argument("--palette", default=None, choices=list(PALETTES.keys()),
                        help="Color palette direction (overrides preset)")
    parser.add_argument("--rendering", default=None, choices=list(RENDERINGS.keys()),
                        help="Visual rendering style (overrides preset)")
    parser.add_argument("--mood", default=None, choices=list(MOODS.keys()),
                        help="Overall mood intensity (overrides preset)")
    parser.add_argument("--layout", default=None, choices=list(LAYOUTS.keys()),
                        help="Information layout (primarily for guide type)")
    parser.add_argument("--text-overlay", default=None,
                        help='Text to render on poster, JSON format: \'{"title":"...", "price":"..."}\'')
    parser.add_argument("--size", default=None,
                        help="Override size as WxH or W*H (default: auto from --type)")
    parser.add_argument("--reference-image", action="append", default=None,
                        help="Path to a reference image for prompt guidance (not sent to API). "
                             "Can be specified multiple times, up to 10.")
    parser.add_argument("--negative-prompt", default=None,
                        help="Additional negative prompt terms")
    parser.add_argument("--poll-interval", type=int, default=2,
                        help="Task poll interval in seconds (default: 2)")
    parser.add_argument("--max-wait", type=int, default=300,
                        help="Max seconds to wait for task completion (default: 300)")
    parser.add_argument("--format", default="text", choices=["text", "json"],
                        help="Output format (default: text)")
    args = parser.parse_args(argv)

    palette, rendering, mood = resolve_style(
        args.style_preset, args.palette, args.rendering, args.mood,
    )

    poster_cfg = POSTER_TYPES[args.type]
    try:
        w, h = parse_size(args.size) if args.size else parse_size(poster_cfg["size"])
    except ValueError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    api_size = size_to_qwen_size(w, h)

    text_overlay = None
    if args.text_overlay:
        try:
            text_overlay = json.loads(args.text_overlay)
        except json.JSONDecodeError as e:
            print(f"Error: --text-overlay is not valid JSON: {e}", file=sys.stderr)
            sys.exit(1)

    ref_paths: List[str] = []
    if args.reference_image:
        if len(args.reference_image) > MAX_REFERENCE_IMAGES:
            print(
                f"Error: too many reference images ({len(args.reference_image)}). "
                f"Max {MAX_REFERENCE_IMAGES}.",
                file=sys.stderr,
            )
            sys.exit(1)
        try:
            ref_paths = validate_reference_image_paths(args.reference_image)
        except FileNotFoundError as e:
            print(f"Error: {e}", file=sys.stderr)
            sys.exit(1)
        print(
            "Warning: Qwen-Image is text-to-image only; reference paths are used for "
            "prompt guidance, not uploaded to the API.",
            file=sys.stderr,
        )

    import sophnet_tools
    api_key = sophnet_tools.get_api_key()
    if not api_key:
        print("Error: No API key found. Set the SOPH_API_KEY environment variable.",
              file=sys.stderr)
        sys.exit(1)

    prompt, negative_prompt = build_prompt(
        args.prompt, w, h,
        palette=palette, rendering=rendering, mood=mood,
        layout=args.layout, text_overlay=text_overlay,
        reference_image_paths=ref_paths or None,
    )
    if args.negative_prompt:
        negative_prompt = f"{negative_prompt}, {args.negative_prompt}"

    print("STATUS=generating", file=sys.stderr)

    try:
        task_id, image_url = generate_image(
            api_key, prompt, api_size,
            negative_prompt=negative_prompt,
            poll_interval=args.poll_interval,
            max_wait=args.max_wait,
        )
    except requests.HTTPError as e:
        handle_http_error(e)
    except requests.RequestException as e:
        print("STATUS=api_error", file=sys.stderr)
        print(f"Error: SophNet Qwen-Image API call failed: {e}", file=sys.stderr)
        sys.exit(1)
    except (TimeoutError, RuntimeError, ValueError) as e:
        print("STATUS=api_error", file=sys.stderr)
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    print(format_output(
        poster_type=args.type, w=w, h=h, api_size=api_size,
        palette=palette, rendering=rendering, mood=mood,
        layout=args.layout, style_preset=args.style_preset,
        status="succeeded", image_url=image_url, task_id=task_id,
        reference_image_count=len(ref_paths),
        fmt=args.format,
    ))


if __name__ == "__main__":
    main()
