# prompt_rewrite.py
import os
from huggingface_hub import InferenceClient

# -------------------------
# Utils
# -------------------------

def get_caption_language(prompt: str) -> str:
    ranges = [
        ('\u4e00', '\u9fff'),  # CJK Unified Ideographs
    ]
    for char in prompt:
        if any(start <= char <= end for start, end in ranges):
            return 'zh'
    return 'en'


def _get_client():
    api_key = os.environ.get("HF_TOKEN") or os.environ.get("hf")
    if not api_key:
        raise EnvironmentError("HF_TOKEN is not set.")
    return InferenceClient(
        provider="cerebras",
        api_key=api_key,
    )


# -------------------------
# Core engine
# -------------------------

def polish_prompt(original_prompt: str, system_prompt: str) -> str:
    client = _get_client()

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": original_prompt},
    ]

    try:
        completion = client.chat.completions.create(
            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
            messages=messages,
        )
        return completion.choices[0].message.content.strip().replace("\n", " ")
    except Exception as e:
        print(f"[prompt_rewrite] Error: {e}")
        return original_prompt


# -------------------------
# System prompts
# -------------------------

SYSTEM_PROMPT_EN = """\
# Image Prompt Rewriting Expert
You are a world-class expert in crafting image prompts for text-to-image models.
Your goal is to accept a short, simple user prompt and rewrite it into a detailed, descriptive, and high-quality prompt that will yield the best aesthetic results.
Focus on:
- Detailed descriptions of the subject, lighting, composition, and texture.
- Adding artistic style descriptors if appropriate (e.g., "cinematic lighting", "photorealistic", "8k resolution").
- Maintaining the original intent of the user's request.
ONLY output the rewritten prompt. Do not add explanations.
"""

SYSTEM_PROMPT_ZH = """\
# 图像 Prompt 改写专家
你是一位世界顶级的图像 Prompt 构建专家，擅长为文生图模型编写提示词。
你的目标是将用户输入的简单提示词，改写为详细、生动且高质量的提示词，以获得最佳的审美效果。
请关注：
- 详细描述主体、光线、构图和纹理。
- 加入通过的艺术风格限定词（例如：“电影质感”、“真实感”、“8k分辨率”）。
- 保持用户原始请求的意图。
仅输出改写后的 Prompt，不要包含任何解释。
"""


# -------------------------
# Public API
# -------------------------

def polish_prompt_en(original_prompt: str) -> str:
    return polish_prompt(original_prompt.strip(), SYSTEM_PROMPT_EN)


def polish_prompt_zh(original_prompt: str) -> str:
    return polish_prompt(original_prompt.strip(), SYSTEM_PROMPT_ZH)


def rewrite(prompt: str) -> str:
    lang = get_caption_language(prompt)
    if lang == "zh":
        return polish_prompt_zh(prompt)
    return polish_prompt_en(prompt)