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

from src.services.provider.base import BaseImageService, ServiceResponse, TaskStatus, GenerationResult
from src.config.settings import MIDJOURNEY_API_KEY, MIDJOURNEY_PROXY_URL, YOUCHUAN_APP_ID, YOUCHUAN_SECRET_KEY

logger = logging.getLogger(__name__)

class MidjourneyImageService(BaseImageService):
    def __init__(self, model_name: str = "midjourney", api_key: Optional[str] = None, api_secret: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key=api_key, **kwargs)
        self.api_secret = api_secret
        
        # Fallback to env vars if not injected
        if not self.api_key:
            self.api_key = YOUCHUAN_APP_ID or MIDJOURNEY_API_KEY
        if not self.api_secret:
            self.api_secret = YOUCHUAN_SECRET_KEY

        # Determine if we use Youchuan (requires both AppID and Secret)
        # 注意: If api_key is set but api_secret is not, it might be Proxy mode (api_key=secret)
        # If both are set, it's Youchuan.
        self.use_youchuan = bool(self.api_key and self.api_secret)
        
        if not self.use_youchuan and not MIDJOURNEY_PROXY_URL:
             # Not Youchuan, we need Proxy URL. 
             # Proxy URL is set, we can try Proxy mode even if api_key is missing (some proxies don't need auth)
            # But usually they do.
             logger.warning("Neither (YOUCHUAN_APP_ID + SECRET) nor MIDJOURNEY_PROXY_URL is set. Midjourney generation will fail.")
            
    async def generate(self, prompt: str, **kwargs) -> ServiceResponse:
        """ 生成 image using Youchuan API or Midjourney Proxy.
        """
        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', {})

        # Get additional keys from extra_config
        api_secret = extra_config.get('api_secret') or extra_config.get('secret_key')

        # Determine mode: Youchuan (app_id + secret) or Proxy (api_key as proxy_url or secret)
        # If api_secret is provided, use Youchuan mode
        # If no api_secret but api_key looks like a URL, use Proxy mode with api_key as proxy_url
        use_youchuan = bool(api_secret)
        proxy_url = None
        if not use_youchuan:
            # Check if api_key is a proxy URL
            if api_key and (api_key.startswith('http://') or api_key.startswith('https://')):
                proxy_url = api_key

        # Pass keys to internal methods via kwargs
        kwargs['_api_key'] = api_key
        kwargs['_api_secret'] = api_secret
        kwargs['_proxy_url'] = proxy_url
        kwargs['_use_youchuan'] = use_youchuan

        if use_youchuan:
            return await self._generate_youchuan(prompt, **kwargs)
        else:
            return await self._generate_proxy(prompt, **kwargs)

    # 类级别批量任务映射
    _batch_tasks: Dict[str, List[str]] = {}

    async def _generate_youchuan_single(self, prompt: str, **kwargs) -> str:
        """生成单张图片，返回 task_id"""
        url = "https://ali.youchuan.cn/v1/tob/diffusion"
        api_key = kwargs.get('_api_key')
        api_secret = kwargs.get('_api_secret')
        headers = {
            "Content-Type": "application/json",
            "x-youchuan-app": api_key,
            "x-youchuan-secret": api_secret
        }
        payload = {
            "text": prompt,
            "callback": kwargs.get("callback", "")
        }

        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(url, headers=headers, json=payload)
            resp.raise_for_status()
            result = resp.json()
        task_id = result.get("taskId") or result.get("data", {}).get("taskId") or result.get("id")

        if not task_id:
            raise ValueError(f"No taskId in Youchuan response: {result}")

        return str(task_id)

    async def _generate_youchuan(self, prompt: str, **kwargs) -> ServiceResponse:
        """生成图片，支持多张（通过多次调用）"""
        # Midjourney/Youchuan 限制最多 4 张
        n = min(kwargs.get("n", 1), 4)

        try:
            task_ids = []
            for i in range(n):
                logger.info(f"Youchuan starting task {i+1}/{n}")
                task_id = await self._generate_youchuan_single(prompt, **kwargs)
                task_ids.append(task_id)

            # 存储批量任务映射
            primary_task_id = task_ids[0] if task_ids else None
            if primary_task_id and len(task_ids) > 1:
                MidjourneyImageService._batch_tasks[primary_task_id] = task_ids

            return ServiceResponse(
                status=TaskStatus.PENDING,
                task_id=primary_task_id,
                meta={
                    "batch_task_ids": task_ids,
                    "batch_size": n,
                    "is_batch": len(task_ids) > 1,
                    "api_key": kwargs.get('_api_key'),
                    "api_secret": kwargs.get('_api_secret')
                }
            )
        except Exception as e:
            logger.error(f"Youchuan Gen Failed: {str(e)}")
            return ServiceResponse(
                task_id="failed",
                status=TaskStatus.FAILED,
                error=str(e)
            )

    async def _generate_proxy_single(self, prompt: str, **kwargs) -> str:
        """生成单张图片，返回 task_id"""
        proxy_url = kwargs.get('_proxy_url') or MIDJOURNEY_PROXY_URL
        if not proxy_url:
            raise ValueError("No proxy URL available for Midjourney")
        base_url = proxy_url.rstrip("/")
        url = f"{base_url}/mj/submit/imagine"

        api_key = kwargs.get('_api_key')
        headers = {"Content-Type": "application/json"}
        if api_key and not (api_key.startswith('http://') or api_key.startswith('https://')):
            headers["mj-api-secret"] = api_key

        payload = {
            "prompt": prompt,
            "base64Array": [],
            "notifyHook": "",
            "state": ""
        }

        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(url, headers=headers, json=payload)
            if response.status_code != 200:
                logger.error(f"MJ Proxy Error: {response.status_code} - {response.text}")
                response.raise_for_status()
            result = response.json()
        task_id = result.get("result") or result.get("id")

        if not task_id:
            raise ValueError(f"No task ID returned from MJ Proxy: {result}")

        return str(task_id)

    async def _generate_proxy(self, prompt: str, **kwargs) -> ServiceResponse:
        """生成图片，支持多张（通过多次调用）"""
        # Midjourney Proxy 限制最多 4 张
        n = min(kwargs.get("n", 1), 4)

        try:
            task_ids = []
            for i in range(n):
                logger.info(f"MJ Proxy starting task {i+1}/{n}")
                task_id = await self._generate_proxy_single(prompt, **kwargs)
                task_ids.append(task_id)

            # 存储批量任务映射
            primary_task_id = task_ids[0] if task_ids else None
            if primary_task_id and len(task_ids) > 1:
                MidjourneyImageService._batch_tasks[primary_task_id] = task_ids

            return ServiceResponse(
                status=TaskStatus.PENDING,
                task_id=primary_task_id,
                meta={
                    "batch_task_ids": task_ids,
                    "batch_size": n,
                    "is_batch": len(task_ids) > 1,
                    "api_key": kwargs.get('_api_key'),
                    "proxy_url": kwargs.get('_proxy_url')
                }
            )

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

    async def check_status(self, task_id: str, **kwargs) -> ServiceResponse:
        # Get keys from kwargs or use system defaults
        api_key = kwargs.get('api_key') or self.api_key
        api_secret = kwargs.get('api_secret') or self.api_secret
        proxy_url = kwargs.get('proxy_url') or MIDJOURNEY_PROXY_URL

        # Determine mode based on available keys
        use_youchuan = bool(api_key and api_secret)

        if use_youchuan:
            return await self._check_status_youchuan(task_id, api_key=api_key, api_secret=api_secret)
        else:
            return await self._check_status_proxy(task_id, api_key=api_key, proxy_url=proxy_url)

    async def _check_status_youchuan_single(self, task_id: str, api_key: str, api_secret: str) -> ServiceResponse:
        """查询单个任务状态"""
        url = "https://ali.youchuan.cn/v1/tob/query"
        headers = {
            "Content-Type": "application/json",
            "x-youchuan-app": api_key,
            "x-youchuan-secret": api_secret
        }
        payload = {"taskId": task_id}

        async with httpx.AsyncClient(timeout=60.0) as client:
            resp = await client.post(url, headers=headers, json=payload)
            data = resp.json()
        inner_data = data.get("data", data)
        status_val = inner_data.get("status")
        img_url = inner_data.get("imgUrl") or inner_data.get("url")

        if str(status_val) in ["2", "SUCCESS", "success", "succeeded"]:
            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                task_id=task_id,
                results=[GenerationResult(url=img_url)] if img_url else None,
                output={"image_url": img_url}
            )
        elif str(status_val) in ["3", "FAILED", "failed", "failure"]:
            return ServiceResponse(
                status=TaskStatus.FAILED,
                task_id=task_id,
                error=f"Youchuan task failed: {data}"
            )
        else:
            return ServiceResponse(
                status=TaskStatus.PROCESSING,
                task_id=task_id
            )

    async def _check_status_youchuan(self, task_id: str, api_key: str, api_secret: str) -> ServiceResponse:
        """查询 Youchuan 任务状态，支持批量"""
        # 检查是否是批量任务
        batch_task_ids = MidjourneyImageService._batch_tasks.get(task_id)

        if not batch_task_ids:
            # 普通单任务
            return await self._check_status_youchuan_single(task_id, api_key, api_secret)

        # 批量任务查询
        logger.info(f"Checking Youchuan batch task {task_id} with {len(batch_task_ids)} sub-tasks")
        all_results = []
        all_succeeded = True
        any_failed = False
        errors = []

        for sub_task_id in batch_task_ids:
            sub_response = await self._check_status_youchuan_single(sub_task_id, api_key, api_secret)

            if sub_response.status == TaskStatus.SUCCEEDED:
                if sub_response.results:
                    all_results.extend(sub_response.results)
            elif sub_response.status == TaskStatus.FAILED:
                any_failed = True
                if sub_response.error:
                    errors.append(sub_response.error)
            else:
                all_succeeded = False

        if all_succeeded and all_results:
            del MidjourneyImageService._batch_tasks[task_id]
            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                task_id=task_id,
                results=all_results,
                meta={"batch_size": len(batch_task_ids), "completed": True}
            )
        elif any_failed and not all_results:
            return ServiceResponse(
                status=TaskStatus.FAILED,
                task_id=task_id,
                error=f"All batch tasks failed: {'; '.join(errors[:3])}",
                meta={"batch_size": len(batch_task_ids), "errors": errors}
            )
        elif any_failed:
            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                task_id=task_id,
                results=all_results,
                meta={"batch_size": len(batch_task_ids), "completed": True, "partial": True}
            )
        else:
            return ServiceResponse(
                status=TaskStatus.PROCESSING,
                task_id=task_id,
                meta={"batch_size": len(batch_task_ids), "completed": False}
            )

    async def _check_status_proxy_single(self, task_id: str, api_key: str, proxy_url: str) -> ServiceResponse:
        """查询单个 MJ Proxy 任务状态"""
        if not proxy_url:
            raise ValueError("MIDJOURNEY_PROXY_URL not set")

        base_url = proxy_url.rstrip("/")
        url = f"{base_url}/mj/task/{task_id}/fetch"

        headers = {}
        if api_key and not (api_key.startswith('http://') or api_key.startswith('https://')):
            headers["mj-api-secret"] = api_key

        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.get(url, headers=headers)
            if response.status_code != 200:
                response.raise_for_status()
            result = response.json()

        status_str = result.get("status")
        status_map = {
            "SUBMITTED": TaskStatus.QUEUED,
            "IN_PROGRESS": TaskStatus.PROCESSING,
            "SUCCESS": TaskStatus.SUCCEEDED,
            "FAILURE": TaskStatus.FAILED
        }
        status = status_map.get(status_str, TaskStatus.UNKNOWN)

        results = []
        if status == TaskStatus.SUCCEEDED:
            image_url = result.get("imageUrl")
            if image_url:
                results.append(GenerationResult(url=image_url, content_type="image/png"))

        return ServiceResponse(
            task_id=result.get("id"),
            status=status,
            output={"image_url": result.get("imageUrl")} if results else {},
            results=results
        )

    async def _check_status_proxy(self, task_id: str, api_key: str, proxy_url: str) -> ServiceResponse:
        """查询 MJ Proxy 任务状态，支持批量"""
        # 检查是否是批量任务
        batch_task_ids = MidjourneyImageService._batch_tasks.get(task_id)

        if not batch_task_ids:
            # 普通单任务
            return await self._check_status_proxy_single(task_id, api_key, proxy_url)

        # 批量任务查询
        logger.info(f"Checking MJ Proxy batch task {task_id} with {len(batch_task_ids)} sub-tasks")
        all_results = []
        all_succeeded = True
        any_failed = False
        errors = []

        for sub_task_id in batch_task_ids:
            sub_response = await self._check_status_proxy_single(sub_task_id, api_key, proxy_url)

            if sub_response.status == TaskStatus.SUCCEEDED:
                if sub_response.results:
                    all_results.extend(sub_response.results)
            elif sub_response.status == TaskStatus.FAILED:
                any_failed = True
                if sub_response.error:
                    errors.append(sub_response.error)
            else:
                all_succeeded = False

        if all_succeeded and all_results:
            del MidjourneyImageService._batch_tasks[task_id]
            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                task_id=task_id,
                results=all_results,
                meta={"batch_size": len(batch_task_ids), "completed": True}
            )
        elif any_failed and not all_results:
            return ServiceResponse(
                status=TaskStatus.FAILED,
                task_id=task_id,
                error=f"All batch tasks failed: {'; '.join(errors[:3])}",
                meta={"batch_size": len(batch_task_ids), "errors": errors}
            )
        elif any_failed:
            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                task_id=task_id,
                results=all_results,
                meta={"batch_size": len(batch_task_ids), "completed": True, "partial": True}
            )
        else:
            return ServiceResponse(
                status=TaskStatus.PROCESSING,
                task_id=task_id,
                meta={"batch_size": len(batch_task_ids), "completed": False}
            )
