"""
连贯生图模块 (Volcengine Doubao)
负责根据分镜脚本和原始素材生成一系列连贯的分镜图片
"""
import base64
import logging
import os
import time
import requests
import json
from pathlib import Path
from typing import List, Dict, Any, Optional
from PIL import Image
import io
from modules import storage

import config
from modules import path_utils

logger = logging.getLogger(__name__)


def _env_int(name: str, default: int) -> int:
    try:
        return int(os.getenv(name, str(default)))
    except Exception:
        return default


# Tunables: slow channels can be hot; default conservative but adjustable.
IMG_SUBMIT_TIMEOUT_S = _env_int("IMG_SUBMIT_TIMEOUT_S", 180)
IMG_POLL_TIMEOUT_S = _env_int("IMG_POLL_TIMEOUT_S", 30)
IMG_MAX_RETRIES = _env_int("IMG_MAX_RETRIES", 3)
IMG_POLL_INTERVAL_S = _env_int("IMG_POLL_INTERVAL_S", 2)
IMG_POLL_MAX_RETRIES = _env_int("IMG_POLL_MAX_RETRIES", 90)  # 90*2s ~= 180s


def _is_retryable_exception(e: Exception) -> bool:
    # Network / transient errors
    if isinstance(e, (requests.Timeout, requests.ConnectionError)):
        return True
    msg = str(e).lower()
    # Transient provider errors often contain these keywords
    if any(k in msg for k in ["timeout", "temporarily", "temporarily unavailable", "gateway", "rate", "try again"]):
        return True
    return False


def _with_retries(fn, *, max_retries: int, label: str):
    last = None
    for attempt in range(1, max_retries + 1):
        try:
            return fn()
        except Exception as e:
            last = e
            retryable = _is_retryable_exception(e)
            logger.warning(f"[{label}] attempt {attempt}/{max_retries} failed: {e} (retryable={retryable})")
            if not retryable or attempt >= max_retries:
                raise
            # small backoff
            time.sleep(min(2 ** (attempt - 1), 4))
    raise last  # pragma: no cover

class ImageGenerator:
    """连贯图片生成器 (Volcengine Provider)"""
    
    def __init__(self):
        self.api_key = config.VOLC_API_KEY
        # Endpoint: https://ark.cn-beijing.volces.com/api/v3/images/generations
        self.endpoint = f"https://ark.cn-beijing.volces.com/api/v3/images/generations"
        self.model = config.IMAGE_MODEL_ID
        
    def _encode_image(self, image_path: str) -> str:
        """读取图片，调整大小并转为 Base64"""
        try:
            with Image.open(image_path) as img:
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                max_size = 1024
                if max(img.size) > max_size:
                    img.thumbnail((max_size, max_size), Image.LANCZOS)
                
                buffer = io.BytesIO()
                img.save(buffer, format="JPEG", quality=80)
                return base64.b64encode(buffer.getvalue()).decode('utf-8')
        except Exception as e:
            logger.error(f"Error processing image {image_path}: {e}")
            return ""

    def generate_single_scene_image(
        self,
        scene: Dict[str, Any],
        original_image_path: Any, 
        previous_image_path: Optional[str] = None,
        model_provider: str = "shubiaobiao", # "shubiaobiao", "gemini", "doubao"
        visual_anchor: str = "",  # 视觉锚点，强制拼接到 prompt 前
        project_id: Optional[str] = None,
    ) -> Optional[str]:
        """
        生成单张分镜图片 (Public)
        """
        scene_id = scene["id"]
        visual_prompt = scene.get("visual_prompt", "")
        
        # 强制拼接 Visual Anchor (确保生图一致性)
        if visual_anchor and visual_anchor not in visual_prompt:
            visual_prompt = f"[{visual_anchor}] {visual_prompt}"
            logger.info(f"Scene {scene_id}: Prepended visual_anchor to prompt")
        
        logger.info(f"Generating image for Scene {scene_id} (Provider: {model_provider})...")
        
        input_images = []
        
        # Handle original_image_path (can be str or list)
        if isinstance(original_image_path, list):
             input_images.extend(original_image_path)
        elif isinstance(original_image_path, str) and original_image_path:
             input_images.append(original_image_path)
             
        if previous_image_path:
            input_images.append(previous_image_path)
            
        try:
            out_dir = path_utils.project_images_dir(project_id) if project_id else config.TEMP_DIR
            out_name = path_utils.unique_filename(
                prefix="scene_image",
                ext="png",
                project_id=project_id,
                scene_id=scene_id,
            )
            output_path = self._generate_single_image(
                prompt=visual_prompt,
                reference_images=input_images, 
                output_filename=out_name,
                provider=model_provider,
                output_dir=out_dir,
            )
            
            if output_path:
                return output_path
            else:
                raise RuntimeError(f"Image generation returned empty for Scene {scene_id}")
                    
        except PermissionError as e:
            logger.error(f"Critical API Error for Scene {scene_id}: {e}")
            raise e
        except Exception as e:
            logger.error(f"Image generation failed for Scene {scene_id}: {e}")
            raise e

    def generate_group_images_doubao(
        self,
        scenes: List[Dict[str, Any]],
        reference_images: List[str],
        visual_anchor: str = "",  # 视觉锚点
        project_id: Optional[str] = None,
    ) -> Dict[int, str]:
        """
        Doubao 组图生成 (Batch) - 拼接 Prompt 一次生成多张
        """
        logger.info("Starting Doubao Group Image Generation...")
        
        # 1. 拼接 Prompts
        # 格式: "Global: [Visual Anchor] ... | S1: ... | S2: ..."
        
        scene_prompts = []
        for scene in scenes:
            # 提取分镜 Visual Prompt
            p = scene.get("visual_prompt", "")
            scene_prompts.append(f"S{scene['id']}:{p}")
            
        combined_scenes_text = " | ".join(scene_prompts)
        
        # 构造 Combined Prompt - 将 visual_anchor 放入 Global 部分
        global_context = f"[{visual_anchor}] Consistent product appearance & style." if visual_anchor else "Consistent product appearance & style."
        combined_prompt = (
            f"Global: {global_context}\n"
            f"{combined_scenes_text}\n"
            "Req: 1 img per scene. Follow specific angles."
        )
        
        logger.info(f"Visual Anchor applied to group prompt: {visual_anchor[:50]}..." if visual_anchor else "No visual_anchor")
        
        # 记录 Prompt 长度供参考
        logger.info(f"Doubao Group Prompt Length: {len(combined_prompt)} chars")

        # 2. 准备 payload
        payload = {
            "model": config.DOUBAO_IMG_MODEL,
            "prompt": combined_prompt,
            "sequential_image_generation": "auto", # 开启组图
            "sequential_image_generation_options": {
                "max_images": len(scenes) # 限制最大张数
            },
            "response_format": "url",
            "size": "1440x2560",
            "stream": False,
            "watermark": False
        }
        
        # 3. 处理参考图
        img_urls = []
        if reference_images:
            for ref_path in reference_images:
                if os.path.exists(ref_path):
                    try:
                        url = storage.upload_file(ref_path)
                        if url: img_urls.append(url)
                    except Exception as e:
                        logger.warning(f"Failed to upload ref image {ref_path}: {e}")
        
        if img_urls:
            payload["image_urls"] = img_urls
            
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {config.VOLC_API_KEY}"
        }
        
        try:
            logger.info(f"Submitting Doubao Group Request (Scenes: {len(scenes)})...")
            resp = requests.post(self.endpoint, json=payload, headers=headers, timeout=240)
            resp.raise_for_status()
            
            data = resp.json()
            results = {}
            
            if "data" in data:
                items = data["data"]
                logger.info(f"Doubao returned {len(items)} images.")
                
                # 尝试将返回的图片映射回 Scene
                # 假设顺序一致
                for i, item in enumerate(items):
                    if i < len(scenes):
                        scene_id = scenes[i]["id"]
                        image_url = item.get("url")
                        
                        if image_url:
                            # Download
                            img_resp = requests.get(image_url, timeout=60)
                            out_dir = path_utils.project_images_dir(project_id) if project_id else config.TEMP_DIR
                            out_name = path_utils.unique_filename(
                                prefix="scene_image",
                                ext="png",
                                project_id=project_id,
                                scene_id=scene_id,
                                extra="group",
                            )
                            output_path = out_dir / out_name
                            with open(output_path, "wb") as f:
                                f.write(img_resp.content)
                            results[scene_id] = str(output_path)
                            
            return results
            
        except Exception as e:
            logger.error(f"Doubao Group Generation Failed: {e}")
            raise e

    def _generate_single_image(
        self, 
        prompt: str, 
        reference_images: List[str],
        output_filename: str,
        provider: str = "shubiaobiao",
        output_dir: Optional[Path] = None,
    ) -> Optional[str]:
        """统一入口"""
        out_dir = output_dir or config.TEMP_DIR
        if provider == "doubao":
            return self._generate_single_image_doubao(prompt, reference_images, output_filename, out_dir)
        elif provider == "gemini":
            return self._generate_single_image_gemini(prompt, reference_images, output_filename, out_dir)
        else:
            return self._generate_single_image_shubiao(prompt, reference_images, output_filename, out_dir)

    def _generate_single_image_doubao(
        self,
        prompt: str,
        reference_images: List[str],
        output_filename: str,
        output_dir: Path
    ) -> Optional[str]:
        """调用 Volcengine Doubao (Image API)"""
        
        # 1. Upload all reference images to R2
        img_urls = []
        if reference_images:
            for ref_path in reference_images:
                if os.path.exists(ref_path):
                    try:
                        url = storage.upload_file(ref_path)
                        if url:
                            img_urls.append(url)
                            logger.info(f"Uploaded Doubao ref image: {url}")
                    except Exception as e:
                        logger.warning(f"Failed to upload Doubao ref image {ref_path}: {e}")
        
        payload = {
            "model": config.DOUBAO_IMG_MODEL, 
            "prompt": prompt,
            "sequential_image_generation": "disabled",
            "response_format": "url",
            "size": "1440x2560", 
            "stream": False,
            "watermark": False
        }
        
        if img_urls:
            payload["image_urls"] = img_urls
            logger.info(f"Doubao Image Payload: prompt='{prompt[:20]}...', image_urls={len(img_urls)}")
        else:
            logger.info(f"Doubao Image Payload: prompt='{prompt[:20]}...', no reference images")

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {config.VOLC_API_KEY}"
        }
        
        def _call():
            logger.info(f"Submitting to Doubao Image: {self.endpoint}")
            resp = requests.post(self.endpoint, json=payload, headers=headers, timeout=IMG_SUBMIT_TIMEOUT_S)
            
            if resp.status_code != 200:
                msg = f"Doubao Image Failed ({resp.status_code}): {resp.text}"
                logger.error(msg)
                raise RuntimeError(msg)
                
            data = resp.json()
            
            if "data" in data and len(data["data"]) > 0:
                image_url = data["data"][0].get("url")
                if image_url:
                    img_resp = requests.get(image_url, timeout=60)
                    img_resp.raise_for_status()
                    
                    output_path = output_dir / output_filename
                    with open(output_path, "wb") as f:
                        f.write(img_resp.content)
                    return str(output_path)
            
            raise RuntimeError(f"No image URL in Doubao response: {data}")
        return _with_retries(_call, max_retries=IMG_MAX_RETRIES, label="doubao_image")

    def _generate_single_image_shubiao(
        self,
        prompt: str,
        reference_images: List[str],
        output_filename: str,
        output_dir: Path
    ) -> Optional[str]:
        """调用 api2img.shubiaobiao.com 通道生成图片（同步返回 base64）"""
        # 准备参考图，内联 base64 方式
        parts = [{"text": prompt}]
        
        # 严格过滤和排序参考图
        valid_refs = []
        if reference_images:
            for p in reference_images:
                if p and os.path.exists(p) and p not in valid_refs:
                    valid_refs.append(p)
        
        logger.info(f"[Shubiaobiao] Input reference images ({len(valid_refs)}): {valid_refs}")
        
        if valid_refs:
            for ref_path in valid_refs:
                try:
                    encoded = self._encode_image(ref_path)
                    if encoded:
                        parts.append({
                            "inlineData": {
                                "mimeType": "image/jpeg",
                                "data": encoded
                            }
                        })
                except Exception as e:
                    logger.error(f"Failed to encode image {ref_path}: {e}")
        
        logger.info(f"[Shubiaobiao] Final payload parts count: {len(parts)} (1 prompt + {len(parts)-1} images)")
        
        payload = {
            "contents": [{
                "role": "user",
                "parts": parts
            }],
            "generationConfig": {
                "responseModalities": ["IMAGE"],
                "imageConfig": {
                    "aspectRatio": "9:16",
                    "imageSize": "2K"
                }
            }
        }
        
        endpoint = f"{config.SHUBIAOBIAO_IMG_BASE_URL}/v1beta/models/{config.SHUBIAOBIAO_IMG_MODEL_NAME}:generateContent"
        headers = {
            "x-goog-api-key": config.SHUBIAOBIAO_IMG_KEY,
            "Content-Type": "application/json"
        }
        
        def _call():
            logger.info(f"Submitting to Shubiaobiao Img: {endpoint}")
            resp = requests.post(endpoint, json=payload, headers=headers, timeout=IMG_SUBMIT_TIMEOUT_S)
            
            if resp.status_code != 200:
                msg = f"Shubiaobiao 提交失败 ({resp.status_code}): {resp.text}"
                logger.error(msg)
                raise RuntimeError(msg)
                
            data = resp.json()
            
            # 查找 base64 图像
            img_b64 = None
            candidates = data.get("candidates") or []
            if candidates:
                content_parts = candidates[0].get("content", {}).get("parts", [])
                for part in content_parts:
                    inline = part.get("inlineData") if isinstance(part, dict) else None
                    if inline and inline.get("data"):
                        img_b64 = inline["data"]
                        break
            
            if not img_b64:
                msg = f"Shubiaobiao 响应缺少图片数据: {data}"
                logger.error(msg)
                raise RuntimeError(msg)
            
            output_path = output_dir / output_filename
            with open(output_path, "wb") as f:
                f.write(base64.b64decode(img_b64))
                
            logger.info(f"Shubiaobiao Generation Success: {output_path}")
            return str(output_path)
        return _with_retries(_call, max_retries=IMG_MAX_RETRIES, label="shubiaobiao_image")

    def _generate_single_image_gemini(
        self,
        prompt: str,
        reference_images: List[str],
        output_filename: str,
        output_dir: Path
    ) -> Optional[str]:
        """调用 Gemini (Wuyin Keji / NanoBanana-Pro) 生成单张图片"""
        
        # 1. 构造 Payload
        payload = {
            "prompt": prompt,
            "aspectRatio": "9:16",
            "imageSize": "2K"
        }
        
        # 处理参考图 (Image-to-Image)
        if reference_images:
            valid_paths = []
            seen = set()
            for p in reference_images:
                if p and os.path.exists(p) and p not in seen:
                    valid_paths.append(p)
                    seen.add(p)
            
            if valid_paths:
                img_urls = []
                for ref_path in valid_paths:
                    try:
                        url = storage.upload_file(ref_path)
                        if url:
                            img_urls.append(url)
                            logger.info(f"Uploaded ref image: {url}")
                    except Exception as e:
                         logger.warning(f"Error uploading ref image {ref_path}: {e}")
                
                if img_urls:
                    payload["img_url"] = img_urls
                    logger.info(f"Using {len(img_urls)} reference images for Gemini Img2Img")

        headers = {
            "Authorization": config.GEMINI_IMG_KEY,
            "Content-Type": "application/json;charset:utf-8"
        }
        
        def _call():
            # 2. 提交任务
            logger.info(f"Submitting to Gemini: {config.GEMINI_IMG_API_URL}")
            resp = requests.post(config.GEMINI_IMG_API_URL, json=payload, headers=headers, timeout=IMG_SUBMIT_TIMEOUT_S)
            
            if resp.status_code != 200:
                msg = f"Gemini 提交失败 ({resp.status_code}): {resp.text}"
                logger.error(msg)
                raise RuntimeError(msg)
                
            data = resp.json()
            if data.get("code") != 200:
                msg = f"Gemini 返回错误: {data}"
                logger.error(msg)
                raise RuntimeError(msg)
                
            task_id = data.get("data", {}).get("id")
            if not task_id:
                raise RuntimeError(f"Gemini 响应缺少 task id: {data}")
                
            logger.info(f"Gemini Task Submitted, ID: {task_id}")
            
            # 3. 轮询状态
            for _ in range(IMG_POLL_MAX_RETRIES):
                time.sleep(IMG_POLL_INTERVAL_S)
                
                poll_url = f"{config.GEMINI_IMG_DETAIL_URL}?key={config.GEMINI_IMG_KEY}&id={task_id}"
                try:
                    poll_resp = requests.get(poll_url, headers=headers, timeout=IMG_POLL_TIMEOUT_S)
                except requests.Timeout:
                    continue
                except Exception as e:
                    continue

                if poll_resp.status_code != 200:
                    continue
                        
                poll_data = poll_resp.json()
                if poll_data.get("code") != 200:
                    raise RuntimeError(f"Gemini 轮询返回错误: {poll_data}")

                result_data = poll_data.get("data", {}) or {}
                status = result_data.get("status") # 0:排队, 1:生成中, 2:成功, 3:失败
                
                if status == 2:
                    image_url = result_data.get("image_url")
                    if not image_url:
                        raise RuntimeError("Gemini 成功但缺少 image_url")
                        
                    logger.info(f"Gemini Generation Success: {image_url}")
                    img_resp = requests.get(image_url, timeout=60)
                    img_resp.raise_for_status()
                    
                    output_path = output_dir / output_filename
                    with open(output_path, "wb") as f:
                        f.write(img_resp.content)
                        
                    return str(output_path)
                    
                if status == 3:
                    fail_reason = result_data.get("fail_reason", "Unknown")
                    raise RuntimeError(f"Gemini 生成失败: {fail_reason}")

            raise RuntimeError("Gemini 生成超时")
        return _with_retries(_call, max_retries=IMG_MAX_RETRIES, label="gemini_image")