import logging
import asyncio
from typing import Optional
from src.services.provider.registry import ModelRegistry, ModelType
from src.services.provider.base import ServiceResponse
from src.services.task_service import task_manager
from src.services.storage_service import storage_manager

from agentscope.tool import ToolResponse
from agentscope.message import TextBlock

logger = logging.getLogger(__name__)

async def _process_and_save_results(task_id: str, results: list, task_type: str = "image") -> list[str]:
    """ 辅助函数 to save results to storage and return new URLs.""" 
    processed_urls = []
    for idx, item in enumerate(results):
        url = item.url
        if not url:
            continue
            
        ext = "png"
        if task_type == "video":
            ext = "mp4"
        
        key = f"generated/{task_type}s/{task_id}_{idx}.{ext}"
        
        try:
            # Save to storage (OSS/Local)
            saved_url = await asyncio.to_thread(storage_manager.save_from_url, url, key)
            processed_urls.append(saved_url or url)
        except Exception as e:
            logger.error(f"Failed to save result to storage: {e}")
            processed_urls.append(url)
            
    return processed_urls

async def generate_image(prompt: str, user_id: Optional[str] = None) -> ToolResponse:
    """ 生成 an image based on the prompt.
    Returns the URL of the generated image.

    Args:
        prompt: Description of the image to generate
        user_id: Current user ID for API key access
    """
    try:
        service = ModelRegistry.get_default(ModelType.IMAGE)
        if not service:
            return ToolResponse(content=[TextBlock(type="text", text="Error: No image generation service available.")])

        # 创建 Task Record
        model_name = getattr(service, "model_id", "default_image")
        task_params = {"prompt": prompt}
        # Run DB operation in thread to avoid blocking event loop
        task_record = await asyncio.to_thread(
            task_manager.create_task,
            type="image",
            model=model_name,
            params=task_params,
            user_id=user_id
        )

        response = await service.generate(prompt=prompt, user_id=user_id)

        # 更新 provider task id if available
        if response.task_id:
            await asyncio.to_thread(
                task_manager.update_task, 
                task_record.id, 
                provider_task_id=response.task_id
            )
        
        # Check if completed immediately (Sync service)
        if response.status == "SUCCEEDED":
             if response.results and len(response.results) > 0:
                 # 进程 and save results
                 final_urls = await _process_and_save_results(task_record.id, response.results, "image")
                 result_data = {"urls": final_urls}
                 
                 await asyncio.to_thread(
                     task_manager.update_task, 
                     task_record.id, 
                     status="success", 
                     result=result_data
                 )
                 return ToolResponse(content=[TextBlock(type="text", text=f"Image generated successfully: {final_urls[0]}")])
             
             await asyncio.to_thread(
                 task_manager.update_task, 
                 task_record.id, 
                 status="failed", 
                 error="Generation succeeded but no results found."
             )
             return ToolResponse(content=[TextBlock(type="text", text="Error: Generation succeeded but no results found.")])

        if not response.task_id:
            error_msg = f"Failed to start generation task. {response.error}"
            await asyncio.to_thread(
                task_manager.update_task, 
                task_record.id, 
                status="failed", 
                error=error_msg
            )
            return ToolResponse(content=[TextBlock(type="text", text=f"Error: {error_msg}")])
        
        task_id = response.task_id
        # Wait loop
        for _ in range(60):
            check_resp = await service.check_status(task_id, user_id=task_record.user_id)
            if check_resp.status == "SUCCEEDED":
                if check_resp.results and len(check_resp.results) > 0:
                    # 进程 and save results
                    final_urls = await _process_and_save_results(task_record.id, check_resp.results, "image")
                    result_data = {"urls": final_urls}
                    
                    await asyncio.to_thread(
                        task_manager.update_task, 
                        task_record.id, 
                        status="success", 
                        result=result_data
                    )
                    return ToolResponse(content=[TextBlock(type="text", text=f"Image generated successfully: {final_urls[0]}")])
                
                await asyncio.to_thread(
                    task_manager.update_task, 
                    task_record.id, 
                    status="failed", 
                    error="Task succeeded but no results found."
                )
                return ToolResponse(content=[TextBlock(type="text", text="Error: Task succeeded but no results found.")])
            elif check_resp.status == "FAILED":
                error_msg = f"Generation failed. {check_resp.error}"
                await asyncio.to_thread(
                    task_manager.update_task, 
                    task_record.id, 
                    status="failed", 
                    error=error_msg
                )
                return ToolResponse(content=[TextBlock(type="text", text=f"Error: {error_msg}")])
            await asyncio.sleep(2)
        
        return ToolResponse(content=[TextBlock(type="text", text=f"Task submitted (ID: {task_id}), but timed out waiting for result.")])
        
    except Exception as e:
        logger.error(f"Tool generate_image error: {e}")
        # Only try to update if task_record exists
        if 'task_record' in locals():
             try:
                 await asyncio.to_thread(
                     task_manager.update_task, 
                     task_record.id, 
                     status="failed", 
                     error=str(e)
                 )
             except Exception:
                 pass
        return ToolResponse(content=[TextBlock(type="text", text=f"Error generating image: {str(e)}")])

async def generate_video(prompt: str, image_url: Optional[str] = None, user_id: Optional[str] = None) -> ToolResponse:
    """ 生成 a video based on the prompt and optional image URL.
    Returns the URL of the generated video.

    Args:
        prompt: Description of the video to generate
        image_url: Optional URL of an image to use as the starting frame
        user_id: Current user ID for API key access
    """
    try:
        service = ModelRegistry.get_default(ModelType.VIDEO)
        if not service:
            return ToolResponse(content=[TextBlock(type="text", text="Error: No video generation service available.")])

        # 创建 Task Record
        model_name = getattr(service, "model_id", "default_video")
        task_params = {"prompt": prompt}
        if image_url:
            task_params["image_url"] = image_url

        task_record = await asyncio.to_thread(
            task_manager.create_task,
            type="video",
            model=model_name,
            params=task_params,
            user_id=user_id
        )

        kwargs = {"user_id": user_id}
        if image_url:
            kwargs["image_url"] = image_url

        response = await service.generate(prompt=prompt, **kwargs)
        
        if response.task_id:
            await asyncio.to_thread(
                task_manager.update_task, 
                task_record.id, 
                provider_task_id=response.task_id
            )
        
        # Check if completed immediately (Sync service)
        if response.status == "SUCCEEDED":
             if response.results and len(response.results) > 0:
                 # 进程 and save results
                 final_urls = await _process_and_save_results(task_record.id, response.results, "video")
                 result_data = {"urls": final_urls}
                 
                 await asyncio.to_thread(
                     task_manager.update_task, 
                     task_record.id, 
                     status="success", 
                     result=result_data
                 )
                 return ToolResponse(content=[TextBlock(type="text", text=f"Video generated successfully: {final_urls[0]}")])
             
             await asyncio.to_thread(
                 task_manager.update_task, 
                 task_record.id, 
                 status="failed", 
                 error="Generation succeeded but no results found."
             )
             return ToolResponse(content=[TextBlock(type="text", text="Error: Generation succeeded but no results found.")])

        if not response.task_id:
            error_msg = f"Failed to start video task. {response.error}"
            await asyncio.to_thread(
                task_manager.update_task, 
                task_record.id, 
                status="failed", 
                error=error_msg
            )
            return ToolResponse(content=[TextBlock(type="text", text=f"Error: {error_msg}")])
        
        task_id = response.task_id
        # Wait loop
        for _ in range(60):
            check_resp = await service.check_status(task_id, user_id=task_record.user_id)
            if check_resp.status == "SUCCEEDED":
                if check_resp.results and len(check_resp.results) > 0:
                    # 进程 and save results
                    final_urls = await _process_and_save_results(task_record.id, check_resp.results, "video")
                    result_data = {"urls": final_urls}
                    
                    await asyncio.to_thread(
                        task_manager.update_task, 
                        task_record.id, 
                        status="success", 
                        result=result_data
                    )
                    return ToolResponse(content=[TextBlock(type="text", text=f"Video generated successfully: {final_urls[0]}")])
                
                await asyncio.to_thread(
                    task_manager.update_task, 
                    task_record.id, 
                    status="failed", 
                    error="Task succeeded but no results found."
                )
                return ToolResponse(content=[TextBlock(type="text", text="Error: Task succeeded but no results found.")])
            elif check_resp.status == "FAILED":
                error_msg = f"Generation failed. {check_resp.error}"
                await asyncio.to_thread(
                    task_manager.update_task, 
                    task_record.id, 
                    status="failed", 
                    error=error_msg
                )
                return ToolResponse(content=[TextBlock(type="text", text=f"Error: {error_msg}")])
            await asyncio.sleep(2)
        return ToolResponse(content=[TextBlock(type="text", text=f"Video generation task submitted (ID: {task_id}). It may take a while to complete.")])

    except Exception as e:
        logger.error(f"Tool generate_video error: {e}")
        if 'task_record' in locals():
             try:
                 await asyncio.to_thread(
                     task_manager.update_task, 
                     task_record.id, 
                     status="failed", 
                     error=str(e)
                 )
             except Exception:
                 pass
        return ToolResponse(content=[TextBlock(type="text", text=f"Error generating video: {str(e)}")])
