"""Image generation interface.

Supports:
- Google Gemini image models via google-genai
- Volcengine Jimeng (即梦) via visual.volcengineapi.com (AK/SK signed)
"""

from typing import Optional, Union, List, Any
from pathlib import Path
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import logging
import time
from ._rate_limit import rate_limited

logger = logging.getLogger(__name__)


DEFAULT_IMAGE_MODEL = "gemini-3-pro-image-preview"
DEFAULT_JIMENG_IMAGE_MODEL = "jimeng_t2i_v40"


class ImageGenerator:
    """Image generation wrapper for multiple providers."""

    def __init__(self, api_key: Optional[str] = None, *, provider: str = "gemini"):
        """Initialize the image generator

        Args:
            api_key: API key for Google Gemini.
            provider: gemini (default) or jimeng
        """
        self.provider = (provider or "gemini").lower().strip()
        self.api_key = api_key
        if self.provider in ("gemini", "google"):
            self.client = genai.Client(api_key=api_key) if api_key else genai.Client()
            self.default_model = DEFAULT_IMAGE_MODEL
        elif self.provider in ("jimeng",):
            self.client = None  # type: ignore[assignment]
            self.default_model = DEFAULT_JIMENG_IMAGE_MODEL
        else:
            raise ValueError("provider must be one of: gemini, google, jimeng")

    def generate_image(
        self,
        prompt: str,
        output_path: Optional[Union[str, Path]] = None,
        model: Optional[str] = None,
        input_images: Optional[List[Union[str, Path, bytes]]] = None,
        response_modalities: Optional[List[str]] = None,
        aspect_ratio: Optional[str] = None,
        image_size: Optional[str] = None,
        output_mime_type: Optional[str] = None,
        output_compression_quality: Optional[int] = None,
        seed: Optional[int] = None,
        width: Optional[int] = None,
        height: Optional[int] = None,
        config: Optional[Union[types.GenerateContentConfig, types.GenerateContentConfigDict]] = None,
        extra_config: Optional[dict[str, Any]] = None,
        use_imagen: bool = True,
        save_description: bool = True
    ) -> tuple[Path, Optional[Path]]:
        """Generate an image based on the prompt and optional input images

        Args:
            prompt: Text prompt for image generation
            output_path: Path to folder where files will be saved, or specific image file.
                        If folder: saves as generated_image.png and generated_image_description.md
                        If file: saves image with that name and description alongside
                        If not provided, defaults to current directory
            model: Model to use for generation. Defaults to gemini-3-pro-image-preview
            input_images: Optional list of image paths to use as reference/context
            response_modalities: Optional response modalities, e.g. ["IMAGE"]
            aspect_ratio: Optional image aspect ratio (e.g. "1:1", "16:9") via image_config
            image_size: Optional image size (e.g. "1K", "2K", "4K") via image_config
            output_mime_type: Optional output mime type (e.g. "image/png", "image/jpeg") via image_config
            output_compression_quality: Optional output compression quality (0-100) via image_config
            config: Optional GenerateContentConfig (merged with explicit args)
            extra_config: Optional dict merged into config (advanced / forward-compatible)
            save_description: Whether to save the prompt as a markdown description file

        Returns:
            tuple[Path, Optional[Path]]: Paths to the saved image and description file (if saved)

        Raises:
            Exception: If image generation fails
        """
        if self.provider == "jimeng":
            return self._generate_image_jimeng(
                prompt,
                output_path,
                model=model,
                seed=seed,
                width=width,
                height=height,
                extra_config=extra_config,
                save_description=save_description,
                # unsupported args (better errors)
                input_images=input_images,
                response_modalities=response_modalities,
                aspect_ratio=aspect_ratio,
                image_size=image_size,
                output_mime_type=output_mime_type,
                output_compression_quality=output_compression_quality,
                config=config,
                use_imagen=use_imagen,
            )

        if not prompt:
            raise ValueError("Prompt cannot be empty")

        if any(v is not None for v in (seed, width, height)):
            raise ValueError("Gemini provider does not support seed/width/height in this API; use image_config or extra_config instead.")

        # Use default model if not specified
        model = model or self.default_model

        # Handle output path - could be folder or file
        if output_path is None:
            output_dir = Path(".")
            image_name = "generated_image"
        else:
            output_path = Path(output_path)
            # Check if it's an existing directory
            if output_path.exists() and output_path.is_dir():
                # It's an existing directory path
                output_dir = output_path
                image_name = "generated_image"
            elif not output_path.suffix:
                # No extension - treat as directory/name combination
                # The parent is the directory, the name is the image name
                if output_path.parent != Path("."):
                    output_dir = output_path.parent
                    image_name = output_path.name
                else:
                    output_dir = Path(".")
                    image_name = output_path.name
            else:
                # It's a file path with extension
                output_dir = output_path.parent if output_path.parent != Path(".") else Path(".")
                image_name = output_path.stem

        # Ensure output directory exists
        output_dir.mkdir(parents=True, exist_ok=True)

        # Set final paths
        image_path = output_dir / f"{image_name}.png"
        description_path = output_dir / f"{image_name}_description.md" if save_description else None

        # Prepare contents list
        contents = []

        # Add input images if provided (support path, bytes, PIL Image)
        if input_images:
            for item in input_images:
                try:
                    if isinstance(item, (str, Path)):
                        img_path = Path(item)
                        if not img_path.exists():
                            raise FileNotFoundError(f"Input image not found: {img_path}")
                        logger.info(f"Loading input image: {img_path}")
                        img = Image.open(img_path)
                        contents.append(img)
                    elif isinstance(item, (bytes, bytearray)):
                        img = Image.open(BytesIO(item))
                        contents.append(img)
                    else:
                        # Try PIL.Image.Image
                        if getattr(item, 'size', None) is not None and hasattr(item, 'mode'):
                            contents.append(item)
                        else:
                            raise TypeError(f"Unsupported input_images item type: {type(item)!r}")
                except Exception as _e:
                    raise

        # Add the text prompt
        contents.append(prompt)

        # Merge config overrides (response_modalities + image_config passthrough).
        cfg_data: dict[str, Any] = {}
        if config is not None:
            if isinstance(config, dict):
                cfg_data.update(config)
            else:
                cfg_data.update(config.model_dump(exclude_none=True))
        if extra_config:
            cfg_data.update(extra_config)
        if response_modalities is not None:
            cfg_data["response_modalities"] = response_modalities

        if any(v is not None for v in (aspect_ratio, image_size, output_mime_type, output_compression_quality)):
            img_cfg: dict[str, Any] = {}
            existing = cfg_data.get("image_config")
            if isinstance(existing, dict):
                img_cfg.update(existing)
            elif existing is not None:
                try:
                    img_cfg.update(existing.model_dump(exclude_none=True))  # type: ignore[attr-defined]
                except Exception:
                    pass
            if aspect_ratio is not None:
                img_cfg["aspect_ratio"] = aspect_ratio
            if image_size is not None:
                img_cfg["image_size"] = image_size
            if output_mime_type is not None:
                img_cfg["output_mime_type"] = output_mime_type
            if output_compression_quality is not None:
                img_cfg["output_compression_quality"] = output_compression_quality
            cfg_data["image_config"] = img_cfg

        cfg = types.GenerateContentConfig.model_validate(cfg_data) if cfg_data else None

        max_retries = 5
        retry_delay = 10  # seconds per requirement

        for attempt in range(max_retries):
            try:
                if attempt > 0:
                    logger.info(f"Retry attempt {attempt + 1}/{max_retries}...")
                    time.sleep(retry_delay)

                if input_images:
                    logger.info(f"Generating image with prompt: {prompt[:100]}... and {len(input_images)} reference image(s)")
                else:
                    logger.info(f"Generating image with prompt: {prompt[:100]}...")

                # Generate content using Gemini API
                response = self.client.models.generate_content(
                    model=model,
                    contents=contents,  # Now includes both images and text
                    config=cfg,
                )

                # Extract image parts from the response
                if response.candidates:
                    image_parts = [
                        part.inline_data.data
                        for part in response.candidates[0].content.parts
                        if part.inline_data
                    ]

                    if image_parts:
                        # Save the first image
                        image = Image.open(BytesIO(image_parts[0]))
                        image.save(image_path)
                        logger.info(f"Image saved to: {image_path}")

                        # Save description if requested
                        if save_description and description_path:
                            description_content = f"# Image Generation Description\n\n"
                            description_content += f"**Generated Image**: {image_name}.png\n\n"
                            description_content += f"## Prompt\n\n{prompt}\n\n"

                            if input_images:
                                description_content += f"## Reference Images\n\n"
                                for img in input_images:
                                    description_content += f"- {Path(img).name}\n"
                                description_content += "\n"

                            description_content += f"## Generation Details\n\n"
                            description_content += f"- **Model**: {model}\n"
                            description_content += f"- **Timestamp**: {time.strftime('%Y-%m-%d %H:%M:%S')}\n"

                            description_path.write_text(description_content, encoding='utf-8')
                            logger.info(f"Description saved to: {description_path}")

                        return image_path, description_path
                    else:
                        # Log text response if any
                        text_parts = [
                            part.text
                            for part in response.candidates[0].content.parts
                            if part.text
                        ]
                        if text_parts:
                            logger.info(f"Model response (text only): {' '.join(text_parts)}")
                        raise Exception("No image was generated in the response")
                else:
                    raise Exception("No response candidates received")

            except Exception as e:
                # Retry on any error
                if attempt < max_retries - 1:
                    logger.warning(f"Error during image generation: {e}. Retrying in {retry_delay}s...")
                    time.sleep(retry_delay)
                    continue
                logger.error(f"Failed to generate image after {max_retries} attempts: {e}")
                raise

        # This should not be reached, but just in case
        raise Exception(f"Failed after {max_retries} attempts")

    def _generate_image_jimeng(
        self,
        prompt: str,
        output_path: Optional[Union[str, Path]] = None,
        *,
        model: Optional[str] = None,
        seed: Optional[int] = None,
        width: Optional[int] = None,
        height: Optional[int] = None,
        extra_config: Optional[dict[str, Any]] = None,
        save_description: bool = True,
        # unsupported args (for clearer validation)
        input_images: Optional[List[Union[str, Path, bytes]]] = None,
        response_modalities: Optional[List[str]] = None,
        aspect_ratio: Optional[str] = None,
        image_size: Optional[str] = None,
        output_mime_type: Optional[str] = None,
        output_compression_quality: Optional[int] = None,
        config: Optional[Union[types.GenerateContentConfig, types.GenerateContentConfigDict]] = None,
        use_imagen: bool = True,
    ) -> tuple[Path, Optional[Path]]:
        from .jimeng_image import JimengImageClient, DEFAULT_JIMENG_IMAGE_MODEL

        if not prompt:
            raise ValueError("Prompt cannot be empty")

        unsupported: dict[str, Any] = {
            "input_images": input_images,
            "response_modalities": response_modalities,
            "aspect_ratio": aspect_ratio,
            "image_size": image_size,
            "output_mime_type": output_mime_type,
            "output_compression_quality": output_compression_quality,
            "config": config,
        }
        bad = [k for k, v in unsupported.items() if v is not None]
        if bad:
            raise ValueError(f"Jimeng provider does not support: {', '.join(bad)}")
        if use_imagen is not True:
            raise ValueError("Jimeng provider does not support use_imagen flag")

        use_model = model or DEFAULT_JIMENG_IMAGE_MODEL
        use_seed = int(seed) if seed is not None else -1

        # Resolve output path (align with existing behavior: always save PNG).
        if output_path is None:
            output_dir = Path(".")
            image_name = "generated_image"
        else:
            output_path = Path(output_path)
            if output_path.exists() and output_path.is_dir():
                output_dir = output_path
                image_name = "generated_image"
            elif not output_path.suffix:
                if output_path.parent != Path("."):
                    output_dir = output_path.parent
                    image_name = output_path.name
                else:
                    output_dir = Path(".")
                    image_name = output_path.name
            else:
                output_dir = output_path.parent if output_path.parent != Path(".") else Path(".")
                image_name = output_path.stem

        output_dir.mkdir(parents=True, exist_ok=True)
        image_path = output_dir / f"{image_name}.png"
        description_path = output_dir / f"{image_name}_description.md" if save_description else None

        client = JimengImageClient()
        task_id = client.submit_t2i(
            prompt=prompt,
            model=use_model,
            seed=use_seed,
            width=width,
            height=height,
            extra=extra_config,
        )
        raw_bytes, result = client.wait_and_download_first_image(model=use_model, task_id=task_id)

        # Normalize to PNG for predictable output.
        img = Image.open(BytesIO(raw_bytes))
        img.save(image_path)

        if description_path is not None:
            content = "# Image Generation Description\n\n"
            content += f"**Provider**: jimeng\n\n"
            content += f"**Model (req_key)**: {use_model}\n\n"
            content += f"**Task ID**: {task_id}\n\n"
            if width is not None or height is not None:
                content += f"**Size**: {width or '?'}x{height or '?'}\n\n"
            status = result.get("status")
            if status:
                content += f"**Status**: {status}\n\n"
            content += f"## Prompt\n\n{prompt}\n"
            description_path.write_text(content, encoding="utf-8")

        return image_path, description_path


@rate_limited(max_calls=8, period=60.0, key="api_sdk.generate_image")
def generate_image(
    prompt: str,
    output_path: Optional[Union[str, Path]] = None,
    api_key: Optional[str] = None,
    provider: str = "gemini",
    model: Optional[str] = None,
    input_images: Optional[List[Union[str, Path, bytes]]] = None,
    response_modalities: Optional[List[str]] = None,
    aspect_ratio: Optional[str] = None,
    image_size: Optional[str] = None,
    output_mime_type: Optional[str] = None,
    output_compression_quality: Optional[int] = None,
    seed: Optional[int] = None,
    width: Optional[int] = None,
    height: Optional[int] = None,
    config: Optional[Union[types.GenerateContentConfig, types.GenerateContentConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
    save_description: bool = True
) -> tuple[Path, Optional[Path]]:
    """Convenience function to generate an image

    Args:
        prompt: Text prompt for image generation
        output_path: Path to folder or file where the image will be saved
        api_key: Optional API key for Google Gemini
        model: Model to use for generation
        input_images: Optional list of image paths to use as reference/context
        response_modalities: Optional response modalities
        aspect_ratio: Optional aspect ratio via image_config
        image_size: Optional image size via image_config
        output_mime_type: Optional output mime type via image_config
        output_compression_quality: Optional output compression quality via image_config
        config: Optional GenerateContentConfig (merged with explicit args)
        extra_config: Optional dict merged into config
        save_description: Whether to save the prompt as a markdown description file

    Returns:
        tuple[Path, Optional[Path]]: Paths to the saved image and description file (if saved)
    """
    generator = ImageGenerator(api_key=api_key, provider=provider)
    return generator.generate_image(
        prompt,
        output_path,
        model,
        input_images,
        response_modalities=response_modalities,
        aspect_ratio=aspect_ratio,
        image_size=image_size,
        output_mime_type=output_mime_type,
        output_compression_quality=output_compression_quality,
        seed=seed,
        width=width,
        height=height,
        config=config,
        extra_config=extra_config,
        save_description=save_description,
    )



@rate_limited(max_calls=8, period=60.0, key="api_sdk.generate_image_bytes")
def generate_image_bytes(
    prompt: str,
    *,
    api_key: Optional[str] = None,
    provider: str = "gemini",
    model: Optional[str] = None,
    input_images: Optional[List[Union[str, Path, bytes]]] = None,
    subject_labels: Optional[List[str]] = None,
    response_modalities: Optional[List[str]] = None,
    aspect_ratio: Optional[str] = None,
    image_size: Optional[str] = None,
    output_mime_type: Optional[str] = None,
    output_compression_quality: Optional[int] = None,
    seed: Optional[int] = None,
    width: Optional[int] = None,
    height: Optional[int] = None,
    config: Optional[Union[types.GenerateContentConfig, types.GenerateContentConfigDict]] = None,
    extra_config: Optional[dict[str, Any]] = None,
) -> bytes:
    """Generate an image and return raw PNG bytes (no file IO).

    Args:
        prompt: Text prompt for image generation
        api_key: Optional API key
        model: Model to use
        input_images: Optional list of image paths or bytes

    Returns:
        Raw PNG bytes for the first generated image
    """
    if (provider or "gemini").lower().strip() == "jimeng":
        from .jimeng_image import JimengImageClient, DEFAULT_JIMENG_IMAGE_MODEL

        if input_images is not None:
            raise ValueError("Jimeng provider does not support input_images in generate_image_bytes")
        if any(v is not None for v in (subject_labels, response_modalities, aspect_ratio, image_size, output_mime_type, output_compression_quality, config)):
            raise ValueError("Jimeng provider in generate_image_bytes only supports: prompt, model(req_key), seed, width, height, extra_config")
        use_model = model or DEFAULT_JIMENG_IMAGE_MODEL
        use_seed = int(seed) if seed is not None else -1

        client = JimengImageClient()
        task_id = client.submit_t2i(prompt=prompt, model=use_model, seed=use_seed, width=width, height=height, extra=extra_config)
        raw_bytes, _ = client.wait_and_download_first_image(model=use_model, task_id=task_id)
        return raw_bytes

    generator = ImageGenerator(api_key=api_key, provider=provider)
    # Compose contents as structured Parts to make subject-image pairing explicit when labels are provided.
    contents: List[object] = []
    if input_images:
        # Normalize to bytes
        norm: List[bytes] = []
        from io import BytesIO
        for item in input_images:
            if isinstance(item, (str, Path)):
                p = Path(item)
                if not p.exists():
                    raise FileNotFoundError(f"Input image not found: {p}")
                with open(p, 'rb') as f:
                    norm.append(f.read())
            elif isinstance(item, (bytes, bytearray)):
                norm.append(bytes(item))
            else:
                # PIL.Image.Image -> convert to PNG bytes
                try:
                    if getattr(item, 'save', None) is not None and getattr(item, 'mode', None) is not None:
                        bio = BytesIO()
                        item.save(bio, format='PNG')
                        norm.append(bio.getvalue())
                    else:
                        raise TypeError
                except Exception:
                    raise TypeError(f"Unsupported input_images item type: {type(item)!r}")
        # Interleave optional labels with image parts
        if subject_labels and len(subject_labels) == len(norm):
            for lbl, data in zip(subject_labels, norm):
                contents.append(types.Part.from_text(text=f"subject: {lbl}"))
                contents.append(types.Part.from_bytes(data=data, mime_type='image/png'))
        else:
            for data in norm:
                contents.append(types.Part.from_bytes(data=data, mime_type='image/png'))
    # Append the main text prompt
    contents.append(types.Part.from_text(text=prompt))
    client = generator.client
    cfg_data: dict[str, Any] = {}
    if config is not None:
        if isinstance(config, dict):
            cfg_data.update(config)
        else:
            cfg_data.update(config.model_dump(exclude_none=True))
    if extra_config:
        cfg_data.update(extra_config)
    if response_modalities is not None:
        cfg_data["response_modalities"] = response_modalities
    if any(v is not None for v in (aspect_ratio, image_size, output_mime_type, output_compression_quality)):
        img_cfg: dict[str, Any] = {}
        existing = cfg_data.get("image_config")
        if isinstance(existing, dict):
            img_cfg.update(existing)
        elif existing is not None:
            try:
                img_cfg.update(existing.model_dump(exclude_none=True))  # type: ignore[attr-defined]
            except Exception:
                pass
        if aspect_ratio is not None:
            img_cfg["aspect_ratio"] = aspect_ratio
        if image_size is not None:
            img_cfg["image_size"] = image_size
        if output_mime_type is not None:
            img_cfg["output_mime_type"] = output_mime_type
        if output_compression_quality is not None:
            img_cfg["output_compression_quality"] = output_compression_quality
        cfg_data["image_config"] = img_cfg
    cfg = types.GenerateContentConfig.model_validate(cfg_data) if cfg_data else None
    max_retries = 5
    retry_delay = 10
    for attempt in range(max_retries):
        try:
            response = client.models.generate_content(model=model or generator.default_model, contents=contents, config=cfg)
            if response.candidates:
                image_parts = [
                    part.inline_data.data
                    for part in response.candidates[0].content.parts
                    if getattr(part, 'inline_data', None)
                ]
                if image_parts:
                    return image_parts[0]
            # No image in response – not a server error; don't retry further here
            raise Exception('No image bytes in response')
        except Exception as e:
            if attempt < max_retries - 1:
                logger.warning(f"Error during image-bytes generation: {e}. Retrying in {retry_delay}s...")
                time.sleep(retry_delay)
                continue
            raise
