# Copyright (C) 2025 AIDC-AI
# Licensed under the MIT License.

import json
import asyncio
import time
from typing import Optional, TypedDict, List, Union
import threading
from collections import defaultdict

from sqlalchemy.orm import identity

from ..utils.globals import set_language, apply_llm_env_defaults
from ..utils.auth_utils import extract_and_store_api_key
import server
from aiohttp import web
import base64

# Import the MCP client function
import os
import shutil

from ..service.debug_agent import debug_workflow_errors
from ..dao.workflow_table import save_workflow_data, get_workflow_data_by_id, update_workflow_ui_by_id
from ..service.mcp_client import comfyui_agent_invoke
from ..utils.request_context import set_request_context, get_session_id
from ..utils.logger import log
from ..utils.modelscope_gateway import ModelScopeGateway
import folder_paths


def get_llm_config_from_headers(request):
    """Extract LLM-related configuration from request headers."""
    return {
        "openai_api_key": request.headers.get('Openai-Api-Key'),
        "openai_base_url": request.headers.get('Openai-Base-Url'),
        # Workflow LLM settings (optional, used by tools/agents that need a different LLM)
        "workflow_llm_api_key": request.headers.get('Workflow-LLM-Api-Key'),
        "workflow_llm_base_url": request.headers.get('Workflow-LLM-Base-Url'),
        "workflow_llm_model": request.headers.get('Workflow-LLM-Model'),
    }


# 全局下载进度存储
download_progress = {}
download_lock = threading.Lock()

# 不再使用内存存储会话消息，改为从前端传递历史消息

# 在文件开头添加
STATIC_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "public")

# Define types using TypedDict
class Node(TypedDict):
    name: str  # 节点名称
    description: str  # 节点描述
    image: str  # 节点图片url，可为空
    github_url: str  # 节点github地址
    from_index: int  # 节点在列表中的位置
    to_index: int  # 节点在列表中的位置

class NodeInfo(TypedDict):
    existing_nodes: List[Node]  # 已安装的节点
    missing_nodes: List[Node]  # 未安装的节点

class Workflow(TypedDict, total=False):
    id: Optional[int]  # 工作流id
    name: Optional[str]  # 工作流名称
    description: Optional[str]  # 工作流描述
    image: Optional[str]  # 工作流图片
    workflow: Optional[str]  # 工作流

class ExtItem(TypedDict):
    type: str  # 扩展类型
    data: Union[dict, list]  # 扩展数据

class ChatResponse(TypedDict):
    session_id: str  # 会话id
    text: Optional[str]  # 返回文本
    finished: bool  # 是否结束
    type: str  # 返回的类型
    format: str  # 返回的格式
    ext: Optional[List[ExtItem]]  # 扩展信息

# 下载进度回调类
class DownloadProgressCallback:
    def __init__(self, id: str, filename: str, file_size: int, download_id: str):
        self.id = id
        self.filename = filename
        self.file_size = file_size
        self.download_id = download_id
        self.progress = 0
        self.status = "downloading"  # downloading, completed, failed
        self.error_message = None
        self.start_time = time.time()
        
        # 初始化进度记录
        with download_lock:
            download_progress[download_id] = {
                "id": id,
                "filename": filename,
                "file_size": file_size,
                "progress": 0,
                "percentage": 0.0,
                "status": "downloading",
                "start_time": self.start_time,
                "estimated_time": None,
                "speed": 0.0,
                "error_message": None
            }

    def update(self, size: int):
        """更新下载进度"""
        self.progress += size
        current_time = time.time()
        elapsed_time = current_time - self.start_time
        
        # 计算下载速度和预估时间
        if elapsed_time > 0:
            speed = self.progress / elapsed_time  # bytes per second
            if speed > 0 and self.progress < self.file_size:
                remaining_bytes = self.file_size - self.progress
                estimated_time = remaining_bytes / speed
            else:
                estimated_time = None
        else:
            speed = 0.0
            estimated_time = None
        
        percentage = (self.progress / self.file_size) * 100 if self.file_size > 0 else 0
        
        # 更新全局进度
        with download_lock:
            if self.download_id in download_progress:
                # 直接更新字典的值，而不是调用update方法
                progress_dict = download_progress[self.download_id]
                progress_dict["progress"] = self.progress
                progress_dict["percentage"] = round(percentage, 2)
                progress_dict["speed"] = round(speed, 2)
                progress_dict["estimated_time"] = round(estimated_time, 2) if estimated_time else None

    def end(self, success: bool = True, error_message: str = None):
        """下载结束回调"""
        current_time = time.time()
        total_time = current_time - self.start_time
        
        if success:
            self.status = "completed"
            # 验证下载完整性
            assert self.progress == self.file_size, f"Download incomplete: {self.progress}/{self.file_size}"
        else:
            self.status = "failed"
            self.error_message = error_message
        
        # 更新最终状态
        with download_lock:
            if self.download_id in download_progress:
                # 直接更新字典的值，而不是调用update方法
                progress_dict = download_progress[self.download_id]
                progress_dict["status"] = self.status
                progress_dict["progress"] = self.file_size if success else self.progress
                if self.file_size > 0:
                    progress_dict["percentage"] = 100.0 if success else (self.progress / self.file_size) * 100
                else:
                    progress_dict["percentage"] = 0.0
                progress_dict["total_time"] = round(total_time, 2)
                progress_dict["error_message"] = self.error_message

    def fail(self, error_message: str):
        """下载失败回调"""
        self.end(success=False, error_message=error_message)

# 生成唯一下载ID
def generate_download_id() -> str:
    import uuid
    return str(uuid.uuid4())

async def upload_to_oss(file_data: bytes, filename: str) -> str:
    # TODO: Implement your OSS upload logic here
    # For now, save locally and return a placeholder URL
    
    try:
        # Create uploads directory if it doesn't exist
        uploads_dir = os.path.join(os.path.dirname(__file__), '..', '..', 'uploads')
        os.makedirs(uploads_dir, exist_ok=True)
        
        # Generate unique filename to avoid conflicts
        import uuid
        unique_filename = f"{uuid.uuid4().hex}_{filename}"
        file_path = os.path.join(uploads_dir, unique_filename)
        
        # Save file locally
        with open(file_path, 'wb') as f:
            f.write(file_data)
        
        # Return a local URL or base64 data URL for now
        # In production, replace this with actual OSS URL
        base64_data = base64.b64encode(file_data).decode('utf-8')
        # Determine MIME type based on file extension
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
            mime_type = f"image/{filename.split('.')[-1].lower()}"
            if mime_type == "image/jpg":
                mime_type = "image/jpeg"
        else:
            mime_type = "image/jpeg"  # Default
            
        return f"data:{mime_type};base64,{base64_data}"
        
    except Exception as e:
        log.error(f"Error uploading file {filename}: {str(e)}")
        # Return original base64 data if upload fails
        base64_data = base64.b64encode(file_data).decode('utf-8')
        return f"data:image/jpeg;base64,{base64_data}"

# 关联用户消息和AI响应的checkpoint信息
def processMessagesWithCheckpoints(messages):
    # ... existing code ...
    pass

@server.PromptServer.instance.routes.post("/api/chat/invoke")
async def invoke_chat(request):
    log.info("Received invoke_chat request")
    
    # Extract and store API key from Authorization header
    extract_and_store_api_key(request)
    
    req_json = await request.json()
    log.info("Request JSON:", req_json)

    response = web.StreamResponse(
        status=200,
        reason='OK',
        headers={
            'Content-Type': 'application/json',
            'X-Content-Type-Options': 'nosniff'
        }
    )
    await response.prepare(request)

    session_id = req_json.get('session_id')
    prompt = req_json.get('prompt')
    images = req_json.get('images', [])
    intent = req_json.get('intent')
    ext = req_json.get('ext')
    historical_messages = req_json.get('messages', [])
    workflow_checkpoint_id = req_json.get('workflow_checkpoint_id')
    
    # 获取当前语言
    language = request.headers.get('Accept-Language', 'en')
    set_language(language)
    
    # 构建配置信息
    config = {
        "session_id": session_id,
        "workflow_checkpoint_id": workflow_checkpoint_id,
        **get_llm_config_from_headers(request),
        "model_select": next((x['data'][0] for x in ext if x['type'] == 'model_select' and x.get('data')), None)
    }
    # Apply .env-based defaults for LLM-related fields (config > .env > code defaults)
    config = apply_llm_env_defaults(config)
    
    # 设置请求上下文 - 这里建立context隔离
    set_request_context(session_id, workflow_checkpoint_id, config)

    # 图片处理已移至前端，图片信息现在包含在历史消息的OpenAI格式中
    # 保留空的处理逻辑以向后兼容，但实际不再使用
    processed_images = []
    if images and len(images) > 0:
        log.info(f"Note: Received {len(images)} images in legacy format, but using OpenAI message format instead")

    # 历史消息已经从前端传递过来，格式为OpenAI格式，直接使用
    log.info(f"-- Received {len(historical_messages)} historical messages")
    
    # Log workflow checkpoint ID if provided (workflow is now pre-saved before invoke)
    if workflow_checkpoint_id:
        log.info(f"Using workflow checkpoint ID: {workflow_checkpoint_id} for session {session_id}")
    else:
        log.info(f"No workflow checkpoint ID provided for session {session_id}")

    # 历史消息已经从前端传递过来，包含了正确格式的OpenAI消息（包括图片）
    # 直接使用前端传递的历史消息，无需重新构建当前消息
    openai_messages = historical_messages
    
    # 不再需要创建用户消息存储到后端，前端负责消息存储

    try:
        # Call the MCP client to get streaming response with historical messages and image support
        # Pass OpenAI-formatted messages and processed images to comfyui_agent_invoke
        accumulated_text = ""
        ext_data = None
        finished = True  # Default to True
        has_sent_response = False
        previous_text_length = 0
        
        log.info(f"config: {config}")
        
        # Pass messages in OpenAI format (images are now included in messages)
        # Config is now available through request context
        async for result in comfyui_agent_invoke(openai_messages, None):
            # The MCP client now returns tuples (text, ext_with_finished) where ext_with_finished includes finished status
            if isinstance(result, tuple) and len(result) == 2:
                text, ext_with_finished = result
                if text:
                    accumulated_text = text  # text from MCP is already accumulated
                    
                if ext_with_finished:
                    # Extract ext data and finished status from the structured response
                    ext_data = ext_with_finished.get("data")
                    finished = ext_with_finished.get("finished", True)
                    log.info(f"-- Received ext data: {ext_data}, finished: {finished}")
            else:
                # Handle single text chunk (backward compatibility)
                text_chunk = result
                if text_chunk:
                    accumulated_text += text_chunk
                    log.info(f"-- Received text chunk: '{text_chunk}', total length: {len(accumulated_text)}")
            
            # Send streaming response if we have new text content
            # Only send intermediate responses during streaming (not the final one)
            if accumulated_text and len(accumulated_text) > previous_text_length:
                chat_response = ChatResponse(
                    session_id=session_id,
                    text=accumulated_text,
                    finished=False,  # Always false during streaming
                    type="message",
                    format="markdown",
                    ext=None  # ext is only sent in final response
                )
                
                await response.write(json.dumps(chat_response).encode() + b"\n")
                previous_text_length = len(accumulated_text)
                await asyncio.sleep(0.01)  # Small delay for streaming effect

        # Send final response with proper finished logic from MCP client
        log.info(f"-- Sending final response: {len(accumulated_text)} chars, ext: {bool(ext_data)}, finished: {finished}")
        
        final_response = ChatResponse(
            session_id=session_id,
            text=accumulated_text,
            finished=finished,  # Use finished status from MCP client
            type="message",
            format="markdown",
            ext=ext_data
        )
        
        await response.write(json.dumps(final_response).encode() + b"\n")

        # AI响应不再存储到后端，前端负责消息存储

    except Exception as e:
        log.error(f"Error in invoke_chat: {str(e)}")
        error_response = ChatResponse(
            session_id=session_id,
            text=f"I apologize, but an error occurred: {str(e)}",
            finished=True,  # Always finish on error
            type="message",
            format="text",
            ext=None
        )
        await response.write(json.dumps(error_response).encode() + b"\n")

    await response.write_eof()
    return response


@server.PromptServer.instance.routes.post("/api/save-workflow-checkpoint")
async def save_workflow_checkpoint(request):
    """
    Save workflow checkpoint for restore functionality
    """
    log.info("Received save-workflow-checkpoint request")
    req_json = await request.json()
    
    try:
        session_id = req_json.get('session_id')
        workflow_api = req_json.get('workflow_api')  # API format workflow
        workflow_ui = req_json.get('workflow_ui')    # UI format workflow  
        checkpoint_type = req_json.get('checkpoint_type', 'debug_start')  # debug_start, debug_complete, user_message_checkpoint
        message_id = req_json.get('message_id')      # User message ID for linking (optional)
        
        if not session_id or not workflow_api:
            return web.json_response({
                "success": False,
                "message": "Missing required parameters: session_id and workflow_api"
            })
        
        # Save workflow with checkpoint type in attributes
        attributes = {
            "checkpoint_type": checkpoint_type,
            "timestamp": time.time()
        }
        
        # Set description and additional attributes based on checkpoint type
        if checkpoint_type == "user_message_checkpoint" and message_id:
            attributes.update({
                "description": f"Workflow checkpoint before user message {message_id}",
                "message_id": message_id,
                "source": "user_message_pre_invoke"
            })
        else:
            attributes["description"] = f"Workflow checkpoint: {checkpoint_type}"
        
        version_id = save_workflow_data(
            session_id=session_id,
            workflow_data=workflow_api,
            workflow_data_ui=workflow_ui,
            attributes=attributes
        )
        
        log.info(f"Workflow checkpoint saved with version ID: {version_id}")
        
        # Return response format based on checkpoint type
        response_data = {
            "version_id": version_id,
            "checkpoint_type": checkpoint_type
        }
        
        if checkpoint_type == "user_message_checkpoint" and message_id:
            response_data.update({
                "checkpoint_id": version_id,  # Add checkpoint_id alias for user message checkpoints
                "message_id": message_id
            })
        
        return web.json_response({
            "success": True,
            "data": response_data,
            "message": f"Workflow checkpoint saved successfully"
        })
        
    except Exception as e:
        log.error(f"Error saving workflow checkpoint: {str(e)}")
        return web.json_response({
            "success": False,
            "message": f"Failed to save workflow checkpoint: {str(e)}"
        })



@server.PromptServer.instance.routes.get("/api/restore-workflow-checkpoint")
async def restore_workflow_checkpoint(request):
    """
    Restore workflow checkpoint by version ID
    """
    log.info("Received restore-workflow-checkpoint request")
    
    try:
        version_id = request.query.get('version_id')
        
        if not version_id:
            return web.json_response({
                "success": False,
                "message": "Missing required parameter: version_id"
            })
        
        try:
            version_id = int(version_id)
        except ValueError:
            return web.json_response({
                "success": False,
                "message": "Invalid version_id format"
            })
        
        # Get workflow data by version ID
        workflow_version = get_workflow_data_by_id(version_id)
        
        if not workflow_version:
            return web.json_response({
                "success": False,
                "message": f"Workflow version {version_id} not found"
            })
        
        log.info(f"Restored workflow checkpoint version ID: {version_id}")
        
        return web.json_response({
            "success": True,
            "data": {
                "version_id": version_id,
                "workflow_data": workflow_version.get('workflow_data'),
                "workflow_data_ui": workflow_version.get('workflow_data_ui'),
                "attributes": workflow_version.get('attributes'),
                "created_at": workflow_version.get('created_at')
            },
            "message": f"Workflow checkpoint restored successfully"
        })
        
    except Exception as e:
        log.error(f"Error restoring workflow checkpoint: {str(e)}")
        return web.json_response({
            "success": False,
            "message": f"Failed to restore workflow checkpoint: {str(e)}"
        })


@server.PromptServer.instance.routes.post("/api/debug-agent")
async def invoke_debug(request):
    """
    Debug agent endpoint for analyzing ComfyUI workflow errors
    """
    log.info("Received debug-agent request")
    
    # Extract and store API key from Authorization header
    extract_and_store_api_key(request)
    
    req_json = await request.json()
    
    response = web.StreamResponse(
        status=200,
        reason='OK',
        headers={
            'Content-Type': 'application/json',
            'X-Content-Type-Options': 'nosniff'
        }
    )
    await response.prepare(request)

    session_id = req_json.get('session_id')
    workflow_data = req_json.get('workflow_data')
    
    # Get configuration from headers (OpenAI settings)
    config = {
        "session_id": session_id,
        "model": "gemini-2.5-flash",  # Default model for debug agents
        **get_llm_config_from_headers(request),
    }
    # Apply .env-based defaults for LLM-related fields (config > .env > code defaults)
    config = apply_llm_env_defaults(config)

    # 获取当前语言
    language = request.headers.get('Accept-Language', 'en')
    set_language(language)
    
    # 设置请求上下文 - 为debug请求建立context隔离
    set_request_context(session_id, None, config)
    
    log.info(f"Debug agent config: {config}")
    log.info(f"Session ID: {session_id}")
    log.info(f"Workflow nodes: {list(workflow_data.keys()) if workflow_data else 'None'}")

    try:
        # Call the debug agent with streaming response
        accumulated_text = ""
        final_ext_data = None
        finished = False
        
        async for result in debug_workflow_errors(workflow_data):
            # Stream the response
            if isinstance(result, tuple) and len(result) == 2:
                text, ext = result
                if text:
                    accumulated_text = text  # text is already accumulated from debug agent
                
                # Handle new ext format from debug_agent matching mcp-client
                if ext:
                    if isinstance(ext, dict) and "data" in ext and "finished" in ext:
                        # New format: {"data": ext, "finished": finished}
                        final_ext_data = ext["data"]
                        finished = ext["finished"]
                        
                        # 检查是否包含需要实时发送的workflow_update或param_update
                        has_realtime_ext = False
                        if final_ext_data:
                            for ext_item in final_ext_data:
                                if ext_item.get("type") in ["workflow_update", "param_update"]:
                                    has_realtime_ext = True
                                    break
                        
                        # 如果包含实时ext数据或者已完成，则发送响应
                        if has_realtime_ext or finished:
                            chat_response = ChatResponse(
                                session_id=session_id,
                                text=accumulated_text,
                                finished=finished,
                                type="message",
                                format="markdown",
                                ext=final_ext_data  # 发送ext数据
                            )
                            await response.write(json.dumps(chat_response).encode() + b"\n")
                        elif not finished:
                            # 只有文本更新，不发送ext数据
                            chat_response = ChatResponse(
                                session_id=session_id,
                                text=accumulated_text,
                                finished=False,
                                type="message",
                                format="markdown",
                                ext=None
                            )
                            await response.write(json.dumps(chat_response).encode() + b"\n")
                    else:
                        # Legacy format: direct ext data (for backward compatibility)
                        final_ext_data = ext
                        finished = False
                        
                        # Create streaming response
                        chat_response = ChatResponse(
                            session_id=session_id,
                            text=accumulated_text,
                            finished=False,
                            type="message",
                            format="markdown",
                            ext=ext
                        )
                        await response.write(json.dumps(chat_response).encode() + b"\n")
                else:
                    # No ext data, just text streaming
                    chat_response = ChatResponse(
                        session_id=session_id,
                        text=accumulated_text,
                        finished=False,
                        type="message",
                        format="markdown",
                        ext=None
                    )
                    await response.write(json.dumps(chat_response).encode() + b"\n")
                
                await asyncio.sleep(0.01)  # Small delay for streaming effect

        # Send final response
        final_response = ChatResponse(
            session_id=session_id,
            text=accumulated_text,
            finished=True,
            type="message",
            format="markdown",
            ext=final_ext_data if final_ext_data else [{"type": "debug_complete", "data": {"status": "completed"}}]
        )
        
        # Save workflow checkpoint after debug completion if we have workflow_data
        if workflow_data and accumulated_text:
            try:
                current_session_id = get_session_id()
                checkpoint_id = save_workflow_data(
                    session_id=current_session_id,
                    workflow_data=workflow_data,
                    workflow_data_ui=None,  # UI format not available in debug agent
                    attributes={
                        "checkpoint_type": "debug_complete",
                        "description": "Workflow state after debug completion",
                        "timestamp": time.time()
                    }
                )
                
                # Add checkpoint info to ext data
                if final_response["ext"]:
                    final_response["ext"].append({
                        "type": "debug_checkpoint",
                        "data": {
                            "checkpoint_id": checkpoint_id,
                            "checkpoint_type": "debug_complete"
                        }
                    })
                else:
                    final_response["ext"] = [{
                        "type": "debug_checkpoint",
                        "data": {
                            "checkpoint_id": checkpoint_id,
                            "checkpoint_type": "debug_complete"
                        }
                    }]
                
                log.info(f"Debug completion checkpoint saved with ID: {checkpoint_id}")
            except Exception as checkpoint_error:
                log.error(f"Failed to save debug completion checkpoint: {checkpoint_error}")
        
        await response.write(json.dumps(final_response).encode() + b"\n")
        log.info("Debug agent processing complete")

    except Exception as e:
        log.error(f"Error in debug agent: {str(e)}")
        import traceback
        traceback.print_exc()
        
        error_response = ChatResponse(
            session_id=session_id,
            text=f"❌ Debug agent error: {str(e)}",
            finished=True,
            type="message",
            format="text",
            ext=[{"type": "error", "data": {"error": str(e)}}]
        )
        await response.write(json.dumps(error_response).encode() + b"\n")

    await response.write_eof()
    return response


@server.PromptServer.instance.routes.post("/api/update-workflow-ui")
async def update_workflow_ui(request):
    """
    Update workflow_data_ui field for a specific checkpoint without affecting other fields
    """
    log.info("Received update-workflow-ui request")
    req_json = await request.json()
    
    try:
        checkpoint_id = req_json.get('checkpoint_id')
        workflow_data_ui = req_json.get('workflow_data_ui')
        
        if not checkpoint_id or not workflow_data_ui:
            return web.json_response({
                "success": False,
                "message": "Missing required parameters: checkpoint_id and workflow_data_ui"
            })
        
        try:
            checkpoint_id = int(checkpoint_id)
        except ValueError:
            return web.json_response({
                "success": False,
                "message": "Invalid checkpoint_id format"
            })
        
        # Update only the workflow_data_ui field
        success = update_workflow_ui_by_id(checkpoint_id, workflow_data_ui)
        
        if success:
            log.info(f"Successfully updated workflow_data_ui for checkpoint ID: {checkpoint_id}")
            return web.json_response({
                "success": True,
                "message": f"Workflow UI data updated successfully for checkpoint {checkpoint_id}"
            })
        else:
            return web.json_response({
                "success": False,
                "message": f"Checkpoint {checkpoint_id} not found"
            })
        
    except Exception as e:
        log.error(f"Error updating workflow UI data: {str(e)}")
        return web.json_response({
            "success": False,
            "message": f"Failed to update workflow UI data: {str(e)}"
        })

@server.PromptServer.instance.routes.post("/api/download-model")
async def download_model(request):
    """
    Download model from ModelScope using SDK
    """
    log.info("Received download-model request")
    req_json = await request.json()
    
    try:
        id = req_json.get('id')
        model_id = req_json.get('model_id')
        model_type = req_json.get('model_type')
        dest_dir = req_json.get('dest_dir')
        
        # 验证必需参数
        if not id:
            return web.json_response({
                "success": False,
                "message": "Missing required parameter: id"
            })

        if not model_id:
            return web.json_response({
                "success": False,
                "message": "Missing required parameter: model_id"
            })
            
        if not model_type:
            return web.json_response({
                "success": False,
                "message": "Missing required parameter: model_type"
            })
        
        log.info(f"Downloading model: {model_id} (type: {model_type})")
        
        # 生成下载ID
        download_id = generate_download_id()
        
        # 计算目标目录：优先使用传入的dest_dir，否则使用ComfyUI的models目录下对应类型
        resolved_dest_dir = None
        if dest_dir:
            resolved_dest_dir = os.path.abspath(os.path.expanduser(f"models/{dest_dir}"))
        else:
            try:
                model_type_paths = folder_paths.get_folder_paths(model_type)
                resolved_dest_dir = model_type_paths[0] if model_type_paths else os.path.join(folder_paths.models_dir, model_type)
            except Exception:
                resolved_dest_dir = os.path.join(folder_paths.models_dir, model_type)

        # 创建进度回调
        progress_callback = DownloadProgressCallback(
            id=id,
            filename=f"{model_id}.{model_type}",
            file_size=0,  # 实际大小会在下载过程中获取
            download_id=download_id
        )

        # 启动下载任务（异步执行）
        async def download_task():
            try:
                # 调用下载方法 - 使用snapshot_download
                from modelscope import snapshot_download
                
                # 创建进度回调包装器 - 实现ModelScope期望的接口
                class ProgressWrapper:
                    """Factory that returns a per-file progress object with update/end."""
                    def __init__(self, download_id: str):
                        self.download_id = download_id

                    def __call__(self, file_name: str, file_size: int):
                        # Create a per-file progress tracker expected by ModelScope
                        download_id = self.download_id

                        class _PerFileProgress:
                            def __init__(self, fname: str, fsize: int):
                                self.file_name = fname
                                self.file_size = max(int(fsize or 0), 0)
                                self.progress = 0
                                self.last_update_time = time.time()
                                self.last_downloaded = 0
                                with download_lock:
                                    if download_id in download_progress:
                                        # If unknown size, keep 0 to avoid div-by-zero
                                        download_progress[download_id]["file_size"] = self.file_size

                            def update(self, size: int):
                                try:
                                    self.progress += int(size or 0)
                                    now = time.time()
                                    # Update global progress
                                    with download_lock:
                                        if download_id in download_progress:
                                            dp = download_progress[download_id]
                                            dp["progress"] = self.progress
                                            if self.file_size > 0:
                                                dp["percentage"] = round(self.progress * 100.0 / self.file_size, 2)
                                            # speed
                                            elapsed = max(now - self.last_update_time, 1e-6)
                                            speed = (self.progress - self.last_downloaded) / elapsed
                                            dp["speed"] = round(speed, 2)
                                    self.last_update_time = now
                                    self.last_downloaded = self.progress
                                except Exception as e:
                                    log.error(f"Error in progress update: {e}")

                            def end(self):
                                # Called by modelscope when a file finishes
                                with download_lock:
                                    if download_id in download_progress:
                                        dp = download_progress[download_id]
                                        if self.file_size > 0:
                                            dp["progress"] = self.file_size
                                            dp["percentage"] = 100.0
                        
                        return _PerFileProgress(file_name, file_size)
 
                progress_wrapper = ProgressWrapper(download_id)
                
                # 添加调试日志
                log.info(f"Starting download with progress wrapper: {download_id}")
                
                # 在线程中执行阻塞的下载，避免阻塞事件循环
                from functools import partial
                local_dir = await asyncio.to_thread(
                    partial(
                        snapshot_download,
                        model_id=model_id,
                        cache_dir=resolved_dest_dir,
                        progress_callbacks=[progress_wrapper]
                    )
                )
                
                # 下载完成
                progress_callback.end(success=True)
                log.info(f"Model downloaded successfully to: {local_dir}")

                # 下载后遍历目录，将所有重要权重/资源文件移动到最外层（与目录同级，即 resolved_dest_dir）
                try:
                    moved_count = 0
                    allowed_exts = {
                        ".safetensors", ".ckpt", ".pt", ".pth", ".bin"
                        # ".msgpack", ".json", ".yaml", ".yml", ".toml", ".png", ".onnx"
                    }
                    for root, dirs, files in os.walk(local_dir):
                        for name in files:
                            ext = os.path.splitext(name)[1].lower()
                            if ext in allowed_exts:
                                src_path = os.path.join(root, name)
                                target_dir = resolved_dest_dir
                                os.makedirs(target_dir, exist_ok=True)
                                target_path = os.path.join(target_dir, name)
                                # 如果已经在目标目录则跳过
                                if os.path.abspath(os.path.dirname(src_path)) == os.path.abspath(target_dir):
                                    continue
                                # 处理重名情况：自动追加 _1, _2 ...
                                if os.path.exists(target_path):
                                    base, ext_real = os.path.splitext(name)
                                    idx = 1
                                    while True:
                                        candidate = f"{base}_{idx}{ext_real}"
                                        candidate_path = os.path.join(target_dir, candidate)
                                        if not os.path.exists(candidate_path):
                                            target_path = candidate_path
                                            break
                                        idx += 1
                                shutil.move(src_path, target_path)
                                moved_count += 1
                    log.info(f"Moved {moved_count} files with extensions {sorted(list(allowed_exts))} to: {resolved_dest_dir}")
                except Exception as move_err:
                    log.error(f"Post-download move failed: {move_err}")
                
            except Exception as e:
                progress_callback.fail(str(e))
                log.error(f"Download failed: {str(e)}")

        # 启动异步下载任务
        asyncio.create_task(download_task())
        
        return web.json_response({
            "success": True,
            "data": {
                "download_id": download_id,
                "id": id,
                "model_id": model_id,
                "model_type": model_type,
                "dest_dir": resolved_dest_dir,
                "status": "started"
            },
            "message": f"Download started for model '{model_id}'"
        })
        
    except ImportError as e:
        log.error(f"ModelScope SDK not installed: {str(e)}")
        return web.json_response({
            "success": False,
            "message": "ModelScope SDK not installed. Please install with: pip install modelscope"
        })
        
    except Exception as e:
        log.error(f"Error starting download: {str(e)}")
        import traceback
        traceback.print_exc()
        return web.json_response({
            "success": False,
            "message": f"Failed to start download: {str(e)}"
        })

@server.PromptServer.instance.routes.get("/api/download-progress/{download_id}")
async def get_download_progress(request):
    """
    Get download progress by download ID
    """
    download_id = request.match_info.get('download_id')
    
    if not download_id:
        return web.json_response({
            "success": False,
            "message": "Missing download_id parameter"
        })
    
    with download_lock:
        progress_info = download_progress.get(download_id)
    
    if not progress_info:
        return web.json_response({
            "success": False,
            "message": f"Download ID {download_id} not found"
        })
    
    return web.json_response({
        "success": True,
        "data": progress_info,
        "message": "Download progress retrieved successfully"
    })

@server.PromptServer.instance.routes.get("/api/download-progress")
async def list_downloads(request):
    """
    List all active downloads
    """
    with download_lock:
        downloads = list(download_progress.keys())
    
    return web.json_response({
        "success": True,
        "data": {
            "downloads": downloads,
            "count": len(downloads)
        },
        "message": "Download list retrieved successfully"
    })

@server.PromptServer.instance.routes.delete("/api/download-progress/{download_id}")
async def clear_download_progress(request):
    """
    Clear download progress record (for cleanup)
    """
    download_id = request.match_info.get('download_id')
    
    if not download_id:
        return web.json_response({
            "success": False,
            "message": "Missing download_id parameter"
        })
    
    with download_lock:
        if download_id in download_progress:
            del download_progress[download_id]
            return web.json_response({
                "success": True,
                "message": f"Download progress {download_id} cleared successfully"
            })
        else:
            return web.json_response({
                "success": False,
                "message": f"Download ID {download_id} not found"
            })

@server.PromptServer.instance.routes.get("/api/model-searchs")
async def model_suggests(request):
    """
    Get model search list by keyword
    """
    log.info("Received model-search request")
    try:
        keyword = request.query.get('keyword')

        if not keyword:
            return web.json_response({
                "success": False,
                "message": "Missing required parameter: keyword"
            })

        # 创建ModelScope网关实例
        gateway = ModelScopeGateway()

        suggests = gateway.search(name=keyword)

        list = suggests["data"] if suggests.get("data") else []

        return web.json_response({
            "success": True,
            "data": {
                "searchs": list,
                "total": len(list)
            },
            "message": f"Get searchs successfully"
        })
        
    except ImportError as e:
        log.error(f"ModelScope SDK not installed: {str(e)}")
        return web.json_response({
            "success": False,
            "message": "ModelScope SDK not installed. Please install with: pip install modelscope"
        })  

    except Exception as e:
        log.error(f"Error get model searchs: {str(e)}")
        import traceback
        traceback.print_exc()
        return web.json_response({
            "success": False,
            "message": f"Get model searchs failed: {str(e)}"
        })
        
@server.PromptServer.instance.routes.get("/api/model-paths")
async def model_paths(request):
    """
    Get model paths by type
    """
    log.info("Received model-paths request")
    try:
        model_paths = list(folder_paths.folder_names_and_paths.keys())
        return web.json_response({
            "success": True,
            "data": {
                "paths": model_paths,
            },
            "message": f"Get paths successfully"
        })
        
    except Exception as e:
        log.error(f"Error get model path: {str(e)}")
        import traceback
        traceback.print_exc()
        return web.json_response({
            "success": False,
            "message": f"Get model failed: {str(e)}"
        })