import os
import httpx
import logging
from typing import Dict, Any, Optional, List

from src.services.provider.base import BaseImageService, ServiceResponse, TaskStatus, GenerationResult
from src.config.settings import MINIMAX_API_KEY

logger = logging.getLogger(__name__)

class MiniMaxImageService(BaseImageService):
    def __init__(self, model_name: str = "image-01", api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key=api_key, **kwargs)

        if not self.api_key:
            self.api_key = MINIMAX_API_KEY

        if not self.api_key:
            logger.warning(f"API Key not found for MiniMax.")

        self.base_url = "https://api.minimaxi.com/v1"

    async def generate(self, prompt: str, **kwargs) -> ServiceResponse:
        """Generate image using MiniMax."""
        user_id = kwargs.get('user_id')
        key_config = self.get_effective_api_key_with_config(user_id)
        if not key_config or not key_config.get('api_key'):
            logger.error(f"No API key available for {self._provider_id}")
            return ServiceResponse(
                status=TaskStatus.FAILED,
                error="API key not configured"
            )

        api_key = key_config['api_key']
        extra_config = key_config.get('extra_config', {})

        url = f"{self.base_url}/image_generation"

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }

        # Base payload
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "n": kwargs.get("n", 1),
            "response_format": "url",
            "prompt_optimizer": kwargs.get("prompt_optimizer", False),
            "aigc_watermark": kwargs.get("aigc_watermark", False)
        }

        # Aspect Ratio
        if kwargs.get("aspect_ratio"):
            payload["aspect_ratio"] = kwargs["aspect_ratio"]
        else:
            payload["aspect_ratio"] = "1:1"

        # Width/Height (Only for image-01)
        if self.model_name == "image-01":
            if kwargs.get("width") and kwargs.get("height"):
                payload["width"] = kwargs["width"]
                payload["height"] = kwargs["height"]

        # Seed
        if kwargs.get("seed"):
            payload["seed"] = int(kwargs["seed"])

        # Style (Only for image-01-live)
        if self.model_name == "image-01-live" and kwargs.get("style"):
            style_input = kwargs["style"]
            if isinstance(style_input, dict):
                 payload["style"] = style_input
            elif isinstance(style_input, str):
                 payload["style"] = {
                     "style_type": style_input,
                     "style_weight": 0.8
                 }

        # Image-to-Image (subject_reference)
        reference_images = kwargs.get("image_inputs") or []

        if reference_images:
            ref_type = kwargs.get("ref_image_type", "character")
            payload["subject_reference"] = [
                {
                    "type": ref_type,
                    "image_file": img_url
                } for img_url in reference_images
            ]

        try:
            logger.info(f"Sending request to MiniMax Image API: {url}, model={self.model_name}")
            async with httpx.AsyncClient(timeout=120.0) as client:
                response = await client.post(url, headers=headers, json=payload)
                response.raise_for_status()
                data = response.json()

            # Check base_resp status
            base_resp = data.get("base_resp", {})
            if base_resp.get("status_code") != 0:
                 error_msg = base_resp.get("status_msg", "Unknown error")
                 raise ValueError(f"MiniMax API Error: {error_msg}")

            # Extract result
            output_data = data.get("data", {})
            image_urls = output_data.get("image_urls", [])
            image_base64 = output_data.get("image_base64", [])

            results = []
            if image_urls:
                for img_url in image_urls:
                    results.append(GenerationResult(
                        url=img_url,
                        orig_prompt=prompt
                    ))
            elif image_base64:
                for b64 in image_base64:
                    results.append(GenerationResult(
                        content=b64,
                        orig_prompt=prompt
                    ))

            if not results:
                raise ValueError("No images returned from MiniMax")

            return ServiceResponse(
                task_id=data.get("id", f"minimax-{os.urandom(4).hex()}"),
                status=TaskStatus.SUCCEEDED,
                results=results,
                meta={
                    "provider_response": data
                }
            )

        except Exception as e:
            logger.error(f"MiniMax Image Generation Failed: {str(e)}")
            return ServiceResponse(
                task_id="failed",
                status=TaskStatus.FAILED,
                error=str(e)
            )

    async def check_status(self, task_id: str) -> ServiceResponse:
        """MiniMax image generation is synchronous."""
        return ServiceResponse(
            status=TaskStatus.UNKNOWN,
            task_id=task_id,
            error="MiniMax image generation is synchronous."
        )
