"""
Canvas Tools Module
Provides canvas manipulation tools for AgentScope agents.
Refactored from canvas_actions.py with enhanced functionality.
""" 
import logging
from typing import List, Dict, Any, Optional
from agentscope.tool import ToolResponse
from agentscope.message import TextBlock

logger = logging.getLogger(__name__)


# ==================== Workflow Creation Tools ====================

async def create_workflow_chain(
    workflow_type: str,
    config: Dict[str, Any]
) -> ToolResponse:
    """ 创建 a complete workflow chain on canvas.
    
    Args:
        workflow_type: Type of workflow. Available types:
            - 'text_to_video': PROMPT_INPUT -> PROMPT_GENERATOR -> IMAGE_GENERATOR -> VIDEO_GENERATOR
            - 'image_enhancement': IMAGE_GENERATOR -> IMAGE_EDITOR -> IMAGE_GENERATOR
            - 'storyboard_generation': INFO_DISPLAY (storyboard) -> IMAGE_GENERATOR (multiple)
            - 'character_showcase': Character info -> Portrait/Three-view -> Video
            - 'scene_generation': Scene info -> Wide/Detail images
        config: Configuration for the workflow (prompts, models, parameters)
    
    Returns:
        ToolResponse with created nodes and connections
    """
    try:
        workflows = {
            'text_to_video': _create_text_to_video_workflow,
            'image_enhancement': _create_image_enhancement_workflow,
            'storyboard_generation': _create_storyboard_workflow,
            'character_showcase': _create_character_showcase_workflow,
            'scene_generation': _create_scene_generation_workflow,
        }
        
        if workflow_type not in workflows:
            return ToolResponse(content=[TextBlock(
                type="text",
                text=f"Unknown workflow type: {workflow_type}. Available: {', '.join(workflows.keys())}"
            )])
        
        result = workflows[workflow_type](config)
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Created {workflow_type} workflow with {len(result['nodes'])} nodes and {len(result['edges'])} connections."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to create workflow: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error creating workflow: {str(e)}"
        )])


async def create_node_group(
    title: str,
    nodes: List[Dict[str, Any]],
    layout: str = "horizontal"
) -> ToolResponse:
    """ 创建 a group of related nodes with automatic layout.
    
    Args:
        title: Group title
        nodes: List of node configurations, each should contain:
            - type: Node type (IMAGE_GENERATOR, VIDEO_GENERATOR, etc.)
            - title: Node title
            - prompt (optional): Node prompt/description
        layout: Layout style ('horizontal', 'vertical', 'grid')
    
    Returns:
        ToolResponse with group and nodes data
    """
    try:
        positioned_nodes = _apply_layout(nodes, layout)
        group_bounds = _calculate_group_bounds(positioned_nodes)
        
        result = {
            "action": "create_node_group",
            "group": {
                "id": f"group_{title.lower().replace(' ', '_')}",
                "title": title,
                **group_bounds
            },
            "nodes": positioned_nodes
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Created group '{title}' with {len(nodes)} nodes in {layout} layout."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to create node group: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error creating node group: {str(e)}"
        )])


# ==================== Node Manipulation Tools ====================

async def update_node_properties(
    node_id: str,
    properties: Dict[str, Any]
) -> ToolResponse:
    """ 更新 properties of an existing node.
    
    Args:
        node_id: Target node ID
        properties: Properties to update (prompt, model, parameters, etc.)
    
    Returns:
        ToolResponse with update confirmation
    """
    try:
        result = {
            "action": "update_node",
            "node_id": node_id,
            "properties": properties
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Updated node {node_id} with {len(properties)} properties."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to update node: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error updating node: {str(e)}"
        )])


async def connect_nodes(
    connections: List[Dict[str, str]]
) -> ToolResponse:
    """ 创建 connections between nodes.
    
    Args:
        connections: List of connection objects, each containing:
            - from: Source node ID
            - to: Target node ID
    
    Returns:
        ToolResponse with connection data
    """
    try:
        result = {
            "action": "create_connections",
            "edges": [
                {
                    "id": f"edge_{conn['from']}_to_{conn['to']}",
                    "source": conn['from'],
                    "target": conn['to']
                }
                for conn in connections
            ]
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Created {len(connections)} connections between nodes."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to connect nodes: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error connecting nodes: {str(e)}"
        )])


async def delete_nodes(
    node_ids: List[str]
) -> ToolResponse:
    """ 删除 nodes from canvas.
    
    Args:
        node_ids: List of node IDs to delete
    
    Returns:
        ToolResponse with deletion confirmation
    """
    try:
        result = {
            "action": "delete_nodes",
            "node_ids": node_ids
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Deleted {len(node_ids)} nodes from canvas."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to delete nodes: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error deleting nodes: {str(e)}"
        )])


# ==================== Canvas Query Tools ====================

async def get_canvas_state() -> ToolResponse:
    """ 获取 current canvas state (nodes, connections, groups).
    
    Returns:
        ToolResponse with canvas state request
    """
    try:
        result = {
            "action": "get_canvas_state",
            "request": True
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text="Requesting current canvas state..."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to get canvas state: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error getting canvas state: {str(e)}"
        )])


async def find_nodes_by_type(
    node_type: str
) -> ToolResponse:
    """ 查找 all nodes of a specific type on canvas.
    
    Args:
        node_type: Node type to search for. Available types:
            - PROMPT_INPUT, PROMPT_GENERATOR, IMAGE_GENERATOR
            - VIDEO_GENERATOR, VIDEO_ANALYZER, AUDIO_GENERATOR
            - INFO_DISPLAY, IMAGE_EDITOR
    
    Returns:
        ToolResponse with search request
    """
    try:
        result = {
            "action": "find_nodes",
            "filter": {"type": node_type}
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Searching for nodes of type: {node_type}"
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to find nodes: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error finding nodes: {str(e)}"
        )])


# ==================== Batch Operations ====================

async def batch_generate(
    node_configs: List[Dict[str, Any]],
    parallel: bool = True
) -> ToolResponse:
    """ 创建 multiple generation nodes.
    
    Args:
        node_configs: List of node configurations, each containing:
            - type: Node type
            - title: Node title
            - prompt: Generation prompt
        parallel: Whether nodes should execute in parallel
    
    Returns:
        ToolResponse with batch operation data
    """
    try:
        result = {
            "action": "batch_generate",
            "nodes": node_configs,
            "parallel": parallel
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Creating {len(node_configs)} generation nodes ({'parallel' if parallel else 'sequential'} execution)."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed batch generation: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error in batch generation: {str(e)}"
        )])


async def apply_template(
    template_name: str,
    variables: Dict[str, Any]
) -> ToolResponse:
    """
    Apply a predefined canvas template with variables.
    
    Args:
        template_name: Name of the template. Available templates:
            - 'character_sheet': Complete character visualization
            - 'scene_breakdown': Scene analysis and generation
            - 'video_production': End-to-end video pipeline
            - 'asset_library': Organized asset display
        variables: Variables to fill in the template
    
    Returns:
        ToolResponse with template application result
    """
    try:
        templates = {
            'character_sheet': _template_character_sheet,
            'scene_breakdown': _template_scene_breakdown,
            'video_production': _template_video_production,
            'asset_library': _template_asset_library,
        }
        
        if template_name not in templates:
            return ToolResponse(content=[TextBlock(
                type="text",
                text=f"Unknown template: {template_name}. Available: {', '.join(templates.keys())}"
            )])
        
        result = templates[template_name](variables)
        result["action"] = "apply_template"
        result["template_name"] = template_name
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Applied template '{template_name}' with {len(result['nodes'])} nodes."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to apply template: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error applying template: {str(e)}"
        )])


# ==================== New Tools ====================

async def auto_layout_nodes(
    node_ids: List[str],
    layout: str = "grid",
    start_position: Optional[Dict[str, int]] = None
) -> ToolResponse:
    """
    Automatically arrange nodes in a specified layout.
    
    Args:
        node_ids: List of node IDs to arrange
        layout: Layout type ('horizontal', 'vertical', 'grid', 'circular')
        start_position: Optional starting position {'x': int, 'y': int}
    
    Returns:
        ToolResponse with new positions for each node
    """
    try:
        start_pos = start_position or {"x": 100, "y": 100}
        
        result = {
            "action": "auto_layout",
            "node_ids": node_ids,
            "layout": layout,
            "start_position": start_pos
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Auto-arranging {len(node_ids)} nodes in {layout} layout."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to auto layout nodes: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error in auto layout: {str(e)}"
        )])


async def duplicate_workflow(
    source_node_ids: List[str],
    offset: Optional[Dict[str, int]] = None
) -> ToolResponse:
    """
    Duplicate a workflow (nodes and their connections).
    
    Args:
        source_node_ids: List of node IDs to duplicate
        offset: Position offset for duplicated nodes {'x': int, 'y': int}
    
    Returns:
        ToolResponse with duplicated workflow data
    """
    try:
        default_offset = offset or {"x": 400, "y": 200}
        
        result = {
            "action": "duplicate_workflow",
            "source_node_ids": source_node_ids,
            "offset": default_offset
        }
        
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Duplicating workflow with {len(source_node_ids)} nodes."
        )], metadata=result)
        
    except Exception as e:
        logger.error(f"Failed to duplicate workflow: {e}")
        return ToolResponse(content=[TextBlock(
            type="text",
            text=f"Error duplicating workflow: {str(e)}"
        )])


# ==================== Helper Functions ====================

def _create_text_to_video_workflow(config: Dict[str, Any]) -> Dict[str, Any]:
    """ 创建 text-to-video workflow."""
    base_x, base_y = 100, 100
    spacing = 300
    
    nodes = [
        {
            "id": "prompt_input",
            "type": "PROMPT_INPUT",
            "title": "输入提示词",
            "prompt": config.get("initial_prompt", ""),
            "position": {"x": base_x, "y": base_y}
        },
        {
            "id": "prompt_gen",
            "type": "PROMPT_GENERATOR",
            "title": "优化提示词",
            "targetType": "image",
            "position": {"x": base_x + spacing, "y": base_y}
        },
        {
            "id": "image_gen",
            "type": "IMAGE_GENERATOR",
            "title": "生成图像",
            "model": config.get("image_model", "flux-1.1-pro"),
            "position": {"x": base_x + spacing * 2, "y": base_y}
        },
        {
            "id": "video_gen",
            "type": "VIDEO_GENERATOR",
            "title": "生成视频",
            "model": config.get("video_model", "kling-v1.6-pro"),
            "position": {"x": base_x + spacing * 3, "y": base_y}
        }
    ]
    
    edges = [
        {"from": "prompt_input", "to": "prompt_gen"},
        {"from": "prompt_gen", "to": "image_gen"},
        {"from": "image_gen", "to": "video_gen"}
    ]
    
    return {"nodes": nodes, "edges": edges}


def _create_image_enhancement_workflow(config: Dict[str, Any]) -> Dict[str, Any]:
    """ 创建 image enhancement workflow."""
    base_x, base_y = 100, 100
    spacing = 300
    
    nodes = [
        {
            "id": "original_image",
            "type": "IMAGE_GENERATOR",
            "title": "原始图像",
            "image": config.get("image_url", ""),
            "position": {"x": base_x, "y": base_y}
        },
        {
            "id": "image_editor",
            "type": "IMAGE_EDITOR",
            "title": "图像编辑",
            "position": {"x": base_x + spacing, "y": base_y}
        },
        {
            "id": "enhanced_image",
            "type": "IMAGE_GENERATOR",
            "title": "增强图像",
            "position": {"x": base_x + spacing * 2, "y": base_y}
        }
    ]
    
    edges = [
        {"from": "original_image", "to": "image_editor"},
        {"from": "image_editor", "to": "enhanced_image"}
    ]
    
    return {"nodes": nodes, "edges": edges}


def _create_storyboard_workflow(config: Dict[str, Any]) -> Dict[str, Any]:
    """ 创建 storyboard generation workflow."""
    storyboards = config.get("storyboards", [])
    base_x, base_y = 100, 100
    spacing_x, spacing_y = 350, 200
    
    nodes = [
        {
            "id": "storyboard_info",
            "type": "INFO_DISPLAY",
            "title": "分镜脚本",
            "content": "\n".join([f"{i+1}. {sb.get('shot_title', '')}" for i, sb in enumerate(storyboards)]),
            "position": {"x": base_x, "y": base_y}
        }
    ]
    
    for i, sb in enumerate(storyboards[:6]):
        row = i // 3
        col = i % 3
        nodes.append({
            "id": f"shot_{i+1}",
            "type": "IMAGE_GENERATOR",
            "title": sb.get("shot_title", f"镜头 {i+1}"),
            "prompt": sb.get("visual_description", ""),
            "position": {
                "x": base_x + spacing_x + col * spacing_x,
                "y": base_y + row * spacing_y
            }
        })
    
    edges = [
        {"from": "storyboard_info", "to": f"shot_{i+1}"}
        for i in range(min(len(storyboards), 6))
    ]
    
    return {"nodes": nodes, "edges": edges}


def _create_character_showcase_workflow(config: Dict[str, Any]) -> Dict[str, Any]:
    """ 创建 character showcase workflow."""
    character = config.get("character", {})
    base_x, base_y = 100, 100
    spacing = 300
    
    nodes = [
        {
            "id": "char_info",
            "type": "INFO_DISPLAY",
            "title": character.get("name", "角色"),
            "content": character.get("desc", ""),
            "position": {"x": base_x, "y": base_y}
        },
        {
            "id": "char_portrait",
            "type": "IMAGE_GENERATOR",
            "title": "角色肖像",
            "prompt": character.get("appearance", ""),
            "template": "character_white_bg",
            "position": {"x": base_x + spacing, "y": base_y - 100}
        },
        {
            "id": "char_three_view",
            "type": "IMAGE_GENERATOR",
            "title": "三视图",
            "prompt": character.get("appearance", ""),
            "template": "character_three_view",
            "position": {"x": base_x + spacing, "y": base_y + 100}
        },
        {
            "id": "char_video",
            "type": "VIDEO_GENERATOR",
            "title": "角色展示视频",
            "template": "asset_360",
            "position": {"x": base_x + spacing * 2, "y": base_y}
        }
    ]
    
    edges = [
        {"from": "char_info", "to": "char_portrait"},
        {"from": "char_info", "to": "char_three_view"},
        {"from": "char_portrait", "to": "char_video"}
    ]
    
    return {"nodes": nodes, "edges": edges}


def _create_scene_generation_workflow(config: Dict[str, Any]) -> Dict[str, Any]:
    """ 创建 scene generation workflow."""
    scene = config.get("scene", {})
    base_x, base_y = 100, 100
    spacing = 300
    
    nodes = [
        {
            "id": "scene_info",
            "type": "INFO_DISPLAY",
            "title": scene.get("name", "场景"),
            "content": f"{scene.get('desc', '')}\n\n位置: {scene.get('location', '')}\n时间: {scene.get('time_of_day', '')}\n氛围: {scene.get('atmosphere', '')}",
            "position": {"x": base_x, "y": base_y}
        },
        {
            "id": "scene_wide",
            "type": "IMAGE_GENERATOR",
            "title": "全景",
            "prompt": f"{scene.get('desc', '')}, wide shot, {scene.get('atmosphere', '')}",
            "position": {"x": base_x + spacing, "y": base_y - 100}
        },
        {
            "id": "scene_detail",
            "type": "IMAGE_GENERATOR",
            "title": "细节",
            "prompt": f"{scene.get('desc', '')}, detailed view, {scene.get('atmosphere', '')}",
            "position": {"x": base_x + spacing, "y": base_y + 100}
        }
    ]
    
    edges = [
        {"from": "scene_info", "to": "scene_wide"},
        {"from": "scene_info", "to": "scene_detail"}
    ]
    
    return {"nodes": nodes, "edges": edges}


def _apply_layout(nodes: List[Dict[str, Any]], layout: str) -> List[Dict[str, Any]]:
    """Apply layout to nodes."""
    base_x, base_y = 100, 100
    spacing_x, spacing_y = 300, 200
    
    positioned = []
    for i, node in enumerate(nodes):
        if layout == "horizontal":
            pos = {"x": base_x + i * spacing_x, "y": base_y}
        elif layout == "vertical":
            pos = {"x": base_x, "y": base_y + i * spacing_y}
        elif layout == "grid":
            cols = 3
            row, col = divmod(i, cols)
            pos = {"x": base_x + col * spacing_x, "y": base_y + row * spacing_y}
        else:
            pos = {"x": base_x, "y": base_y}
        
        positioned.append({**node, "position": pos})
    
    return positioned


def _calculate_group_bounds(nodes: List[Dict[str, Any]]) -> Dict[str, float]:
    """Calculate bounding box for a group of nodes."""
    if not nodes:
        return {"x": 0, "y": 0, "width": 300, "height": 200}
    
    positions = [n["position"] for n in nodes]
    min_x = min(p["x"] for p in positions) - 50
    min_y = min(p["y"] for p in positions) - 50
    max_x = max(p["x"] for p in positions) + 250
    max_y = max(p["y"] for p in positions) + 150
    
    return {
        "x": min_x,
        "y": min_y,
        "width": max_x - min_x,
        "height": max_y - min_y
    }


# 模板 functions
def _template_character_sheet(variables: Dict[str, Any]) -> Dict[str, Any]:
    """Character sheet template.""" 
    return _create_character_showcase_workflow({"character": variables})


def _template_scene_breakdown(variables: Dict[str, Any]) -> Dict[str, Any]:
    """Scene breakdown template.""" 
    return _create_scene_generation_workflow({"scene": variables})


def _template_video_production(variables: Dict[str, Any]) -> Dict[str, Any]:
    """Video production template.""" 
    return _create_text_to_video_workflow(variables)


def _template_asset_library(variables: Dict[str, Any]) -> Dict[str, Any]:
    """Asset library template."""
    assets = variables.get("assets", [])
    base_x, base_y = 100, 100
    spacing_x, spacing_y = 300, 200
    
    nodes = []
    for i, asset in enumerate(assets):
        row = i // 4
        col = i % 4
        nodes.append({
            "id": f"asset_{i}",
            "type": "INFO_DISPLAY",
            "title": asset.get("name", f"资产 {i+1}"),
            "content": asset.get("desc", ""),
            "position": {
                "x": base_x + col * spacing_x,
                "y": base_y + row * spacing_y
            }
        })
    
    return {"nodes": nodes, "edges": []}
