# -*- coding: utf-8 -*-
"""
AI 图片生成器 - Flask 后端
使用 Pollinations AI 免费接口，一句话生成图片/图标/Logo
"""

import urllib.parse

import requests
from flask import Flask, jsonify, render_template, request

app = Flask(__name__)

# Pollinations AI 免费图像生成接口
POLLINATIONS_URL = "https://image.pollinations.ai/prompt/{prompt}"

# 图片类型对应的提示词前缀
TYPE_PREFIX = {
    "image": "a high quality image of",
    "icon": "a clean flat icon of",
    "logo": "a minimalist logo design of",
}

# 风格对应的提示词后缀
STYLE_SUFFIX = {
    "realistic": "realistic, detailed, 4k, photographic",
    "cartoon": "cartoon style, colorful, cute, flat illustration",
    "watercolor": "watercolor painting, soft colors, artistic",
    "cyberpunk": "cyberpunk style, neon lights, futuristic, sci-fi",
    "minimal": "minimalist, clean, simple, flat design, white background",
}

# 尺寸映射（宽 x 高）
SIZE_MAP = {
    "square": (1024, 1024),
    "landscape": (1280, 768),
    "portrait": (768, 1280),
}


@app.route("/")
def index():
    """渲染首页"""
    return render_template("index.html")


@app.route("/api/generate", methods=["POST"])
def generate():
    """生成图片接口"""
    data = request.get_json(silent=True) or {}
    prompt = (data.get("prompt") or "").strip()
    img_type = data.get("type", "image")
    size = data.get("size", "square")
    style = data.get("style", "realistic")

    # 参数校验
    if not prompt:
        return jsonify({"success": False, "error": "请输入图片描述"}), 400

    if img_type not in TYPE_PREFIX:
        img_type = "image"
    if size not in SIZE_MAP:
        size = "square"
    if style not in STYLE_SUFFIX:
        style = "realistic"

    # 组装完整提示词
    full_prompt = f"{TYPE_PREFIX[img_type]} {prompt}, {STYLE_SUFFIX[style]}"

    width, height = SIZE_MAP[size]
    encoded_prompt = urllib.parse.quote(full_prompt)
    image_url = (
        f"{POLLINATIONS_URL.format(prompt=encoded_prompt)}"
        f"?width={width}&height={height}&nologo=true&seed={_random_seed()}"
    )

    # 预检图片是否可访问（Pollinations 有时会延迟生成）
    try:
        head = requests.head(image_url, timeout=30, allow_redirects=True)
        if head.status_code >= 400:
            return jsonify({
                "success": False,
                "error": f"图像服务返回异常状态：{head.status_code}，请稍后重试",
            }), 502
    except requests.RequestException as exc:
        return jsonify({
            "success": False,
            "error": f"无法连接图像服务：{exc}",
        }), 502

    return jsonify({
        "success": True,
        "image_url": image_url,
        "prompt": full_prompt,
    })


def _random_seed() -> int:
    """生成随机种子，避免相同提示词返回相同图片"""
    import random
    return random.randint(1, 10_000_000)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)