import argparse
import re
import sys
import time

from openai import OpenAI

REWRITE_SYSTEM_PROMPT = """你是专业的AI图像生成Prompt工程师的Prompt Engineering Engine,也是一名拥有百科知识和视觉导演能力的创意总监.你的任务是分析用户的原始图像需求,推理出隐含知识和最佳视觉方案,并改写成一个**明确,详细,可直接用于图像生成的英文prompt**.

## 核心目标

图像生成模型只能执行直接的视觉描述,不能自行补全背景知识,逻辑关系或文字内容.因此,你必须提前完成知识解析,空间规划和视觉导演,把结果显式写入prompt中.

使用 SCALIST 框架扩写每个画面:
- **Subject**: 主体的身份,外观,颜色,材质,纹理,动作,表情,服饰.
- **Composition**: 镜头景别,视角,主体位置,前景/中景/背景层次,留白和视觉焦点.
- **Action**: 主体正在做什么,动作方向,姿态,互动关系.
- **Location**: 场景地点,室内/室外,时代,天气,时间段,环境细节.
- **Image style**: photorealistic, cinematic, oil painting, watercolor, anime, 3D render 等,并匹配合适的光线和色彩氛围.
- **Specs**: 摄影/渲染参数,如 85mm lens, low-angle shot, shallow depth of field, soft diffused light, dramatic backlighting, matte texture, sharp focus.
- **Text rendering**: 如果用户要求文字,必须把准确文字放在英文双引号中,并说明字体风格,颜色,大小,材质和精确位置.

1. **知识解析与显式化**: 凡是诗词,歌词,名言,公式,历史人物,科学概念,地标,名画,文化符号,历史事件,UI布局或现实世界对象,都要先解析出具体答案和可见特征,再写入prompt.不要只写 "Mona Lisa","Dunkirk evacuation","freedom" 这类需要模型自行理解的词.
2. **空间与逻辑锚定**: 把模糊关系改写为明确布局,例如 top left corner, centered in the foreground, slightly behind the main subject, background out of focus, text aligned along the bottom edge.不要使用"旁边""一些""好看"等含糊表达.
3. **文字排版精度**: 中文,英文,公式,多语言文本都必须逐字保留在引号中,例如 "床前明月光,疑是地上霜.举头望明月,低头思故乡." 或 "E = mc²";同时指定字体(calligraphy, serif, sans-serif, handwritten),颜色,材质和位置.
4. **真实世界落地**: 如果用户要求事实准确的内容,例如历史文物,天气现象,人物肖像,建筑,仪表盘或应用界面,要使用你的内部知识补全准确视觉细节.
5. **抽象概念具象化**: 把"自由,孤独,未来感,治愈"等抽象词转成可见场景,符号和氛围,例如飞鸟,断裂锁链,辽阔天空,冷色霓虹,柔和晨光等.

## 示例合并学习

- 用户说"李白的静夜思写在墙上",prompt 应写出完整中文诗句,并指定它以优雅中国书法写在古旧石墙的哪个位置.
- 用户说"三大力学的奠基人"或"爱因斯坦写质能方程",prompt 应解析出 Isaac Newton 或 Albert Einstein,并描述人物外貌,时代服饰,黑板,公式 "E = mc²" 等可见内容.
- 用户说"蒙娜丽莎""比萨斜塔""福字""敦刻尔克大撤退",prompt 应描述对应画面特征: 神秘微笑与交叠双手,倾斜白色大理石钟楼与拱廊,红底金色/黑色书法 "福",1940年海滩上等待撤离的士兵和海面船只.

## 输出prompt要求

- prompt 必须是一个英文的,连贯自然的单段落,像 Creative Director's Brief,而不是关键词堆砌或 tag soup.
- 简单需求可以更短,复杂画面可以更长.
- 最重要的主体和画面意图放在开头,然后自然展开构图,动作,地点,风格,技术参数和文字渲染.
- 使用完整句子,丰富但准确的形容词,摄影/绘画/设计术语.
- 不要包含任何需要图像模型继续推理才能理解的表达.
- prompt 必须自包含,仅凭prompt本身就能准确生成图片.

只输出改写后的prompt, 不添加其他任何内容:"""

DEFAULT_MODEL_ID = "HiDream-ai/Prompt-Refine"
DEFAULT_BASE_URL = "http://localhost:8000/v1"

def _has_tail_repetition(caption):
    N = 5
    sentence = caption.lower().split()
    if len(sentence) < 6:
        return True
    grams = [tuple(sentence[i: i + N]) for i in range(len(sentence) - N + 1)]
    grams_dict = {}
    for ent in grams:
        grams_dict[ent] = grams_dict.get(ent, 0) + 1
        if grams_dict[ent] > 11:
            return True
    return False

def _call_with_retry(
    client: OpenAI,
    *,
    model_id: str,
    prompt: str,
    max_tokens: int,
    max_attempts: int = 5,
    retry_delay: float = 30.0,
) -> str | None:
    """Call the chat-completions endpoint, retrying on any exception.

    Returns the assistant message content on success, or ``None`` if every
    attempt raised. Each retry waits ``retry_delay`` seconds before firing.
    """
    for attempt in range(1, max_attempts + 1):
        try:
            resp = client.chat.completions.create(
                model=model_id,
                messages=[
                    {"role": "system", "content": REWRITE_SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=max_tokens,
            )
            return resp.choices[0].message.content
        except Exception as exc:  # noqa: BLE001 - retry on any API/network error
            print(
                f"[refine_prompt] API call attempt {attempt}/{max_attempts} "
                f"failed: {type(exc).__name__}: {exc}",
                file=sys.stderr,
            )
            if attempt < max_attempts:
                print(
                    f"[refine_prompt] sleeping {retry_delay:g}s before retrying...",
                    file=sys.stderr,
                )
                time.sleep(retry_delay)
    return None


def refine_prompt(
    prompt: str,
    model_id: str = DEFAULT_MODEL_ID,
    base_url: str = DEFAULT_BASE_URL,
    api_key: str = "EMPTY",
    max_tokens: int = 2048,
    max_attempts: int = 3,
    api_max_attempts: int = 5,
    api_retry_delay: float = 30.0,
) -> str:
    """Rewrite a raw user prompt via an OpenAI-compatible endpoint.

    The endpoint is expected to serve `HiDream-ai/Prompt-Refine`
    (see `start_vllm_server.sh`).

    Retries up to ``max_attempts`` times if the refined output ends in a
    degenerate word loop. Each refine attempt itself retries the underlying API
    call up to ``api_max_attempts`` times (sleeping ``api_retry_delay`` seconds
    between failures) so transient errors / rate limits don't kill the run.
    If every attempt is degenerate or the API is unreachable, falls back to the
    original ``prompt`` so downstream image generation still gets a usable input.
    """
    client = OpenAI(api_key=api_key, base_url=base_url)
    for attempt in range(1, max_attempts + 1):
        refined = _call_with_retry(
            client,
            model_id=model_id,
            prompt=prompt,
            max_tokens=max_tokens,
            max_attempts=api_max_attempts,
            retry_delay=api_retry_delay,
        )
        if refined is None:
            print(
                f"[refine_prompt] attempt {attempt}/{max_attempts} exhausted API "
                "retries; trying again...",
                file=sys.stderr,
            )
            continue
        if not _has_tail_repetition(refined):
            return refined
        print(
            f"[refine_prompt] attempt {attempt}/{max_attempts} produced "
            "degenerate tail repetition; retrying...",
            file=sys.stderr,
        )
    print(
        f"[refine_prompt] all {max_attempts} attempts failed; falling back to "
        "the original prompt.",
        file=sys.stderr,
    )
    return prompt


def main():
    p = argparse.ArgumentParser("Prompt rewriting agent for HiDream-O1-Image")
    p.add_argument("--prompt", type=str, required=True, help="Raw user prompt to rewrite.")
    p.add_argument("--model_id", type=str, default=DEFAULT_MODEL_ID, help="Model name")
    p.add_argument("--base_url", type=str, default=DEFAULT_BASE_URL,
                   help="OpenAI-compatible base URL")
    args = p.parse_args()
    print(refine_prompt(args.prompt, model_id=args.model_id, base_url=args.base_url))


if __name__ == "__main__":
    main()