import requests
import urllib.parse
import datetime
import random
import os
from io import BytesIO

from  auto_media_publisher.config.conf_base import (
    TRANSLATE_API_TEMPLATE,
    POLLINATIONS_API_TEMPLATE,
    POLLINATIONS_DEFAULT_WIDTH ,
    POLLINATIONS_DEFAULT_HEIGHT ,
    POLLINATIONS_DEFAULT_MODEL,
)


class Text2ImageGenerator:
    def __init__(self, width=POLLINATIONS_DEFAULT_WIDTH, height=POLLINATIONS_DEFAULT_HEIGHT, model=POLLINATIONS_DEFAULT_MODEL):
        self.width = width
        self.height = height
        self.model = model

    def translate_prompt(self, prompt: str) -> str:
        """
        使用 MyMemory API 将中文提示词翻译为英文，失败则返回原文。
        """
        
        encoded_prompt = urllib.parse.quote(prompt)
        api_url = TRANSLATE_API_TEMPLATE.format(prompt=encoded_prompt)
        try:
            res = requests.get(api_url, timeout=30)
            res.raise_for_status()
            data = res.json()
            if data.get("responseStatus") == 200 and "translatedText" in data.get("responseData", {}):
                translated = data["responseData"]["translatedText"]
                print(f"[✓] 翻译成功: {translated}")
                return translated
        except Exception as e:
            print(f"[!] 翻译失败，使用原文: {e}")
        return prompt

    def build_url(self, prompt: str) -> str:
        """
        构建 Pollinations 图片生成 URL。
        """
        seed = random.randint(1, 1_000_000)
        encoded_prompt = urllib.parse.quote(prompt)
        return POLLINATIONS_API_TEMPLATE.format(
            prompt=encoded_prompt,
            width=self.width,
            height=self.height,
            seed=seed,
            model=self.model
        )
    
    def _fetch_with_retry(self, url: str, max_retries: int = 3):
        """
        最多尝试 max_retries 次请求 URL，成功则返回 response.content，否则返回 None。
        """
        for attempt in range(1, max_retries + 1):
            try:
                response = requests.get(url, timeout=30)
                response.raise_for_status()
                return response.content
            except Exception as e:
                print(f"[!] 请求第 {attempt} 次失败：{e}")
        return None

    def generate_image(self, prompt: str, filename: str = None) -> str:
        if not prompt.strip():
            raise ValueError("提示词不能为空")

        translated_prompt = self.translate_prompt(prompt)
        url = self.build_url(translated_prompt)

        content = self._fetch_with_retry(url)
        if content is None:
            return None

        if not filename:
            now = datetime.datetime.now()
            filename = now.strftime("%Y%m%d%H%M%S") + ".png"
        if not filename.lower().endswith(".png"):
            filename += ".png"

        with open(filename, "wb") as f:
            f.write(content)

        abs_path = os.path.abspath(filename)
        print(f"[✓] 图片已保存：{abs_path}")
        return abs_path
    
    def generate_image_to_stream(self, prompt: str) -> BytesIO:
        if not prompt.strip():
            raise ValueError("提示词不能为空")

        translated_prompt = self.translate_prompt(prompt)
        url = self.build_url(translated_prompt)

        content = self._fetch_with_retry(url)
        if content is None:
            return None

        image_stream = BytesIO(content)
        return image_stream