# Generated by Copilot
# Maintained by Dinesh Kumar C (https://github.com/Dinesh210805/blender-mcp)
# Forked from original work by Siddharth Ahuja (https://github.com/ahujasid/blender-mcp)

"""PolyHaven asset tools for Blender MCP."""

from typing import Optional
from mcp.server.fastmcp import Context

from .base import BaseTool, tool_wrapper
from ..utils import validate_asset_params, validate_object_name
from ..exceptions import ValidationError


class PolyHavenStatusTool(BaseTool):
    """Check PolyHaven integration status."""
    
    @tool_wrapper
    def execute(self) -> str:
        """Check if PolyHaven integration is enabled in Blender."""
        result = self.send_command("get_polyhaven_status")
        enabled = result.get("enabled", False)
        message = result.get("message", "")
        
        if enabled:
            message += "\nPolyHaven is good for textures and has a wide variety of assets."
        
        return message


class PolyHavenCategoriesTool(BaseTool):
    """Get PolyHaven asset categories."""
    
    @tool_wrapper
    def execute(self, asset_type: str = "hdris") -> str:
        """Get a list of categories for a specific asset type on PolyHaven."""
        # Validate asset type
        valid_types = ["hdris", "textures", "models", "all"]
        if asset_type not in valid_types:
            raise ValidationError(f"Asset type must be one of: {', '.join(valid_types)}")
        
        result = self.send_command("get_polyhaven_categories", {"asset_type": asset_type})
        
        if "error" in result:
            return f"Error: {result['error']}"
        
        # Format the categories
        categories = result.get("categories", {})
        if not categories:
            return f"No categories found for {asset_type}"
        
        formatted_output = f"Categories for {asset_type}:\\n\\n"
        
        # Sort categories by count (descending)
        sorted_categories = sorted(categories.items(), key=lambda x: x[1], reverse=True)
        
        for category, count in sorted_categories:
            formatted_output += f"- {category}: {count} assets\\n"
        
        return formatted_output


class PolyHavenSearchTool(BaseTool):
    """Search PolyHaven assets."""
    
    @tool_wrapper
    def execute(self, asset_type: str = "all", categories: Optional[str] = None) -> str:
        """Search for assets on PolyHaven with optional filtering."""
        # Validate asset type
        valid_types = ["hdris", "textures", "models", "all"]
        if asset_type not in valid_types:
            raise ValidationError(f"Asset type must be one of: {', '.join(valid_types)}")
        
        result = self.send_command("search_polyhaven_assets", {
            "asset_type": asset_type,
            "categories": categories
        })
        
        if "error" in result:
            return f"Error: {result['error']}"
        
        # Format the results
        assets = result.get("assets", {})
        total_count = result.get("total_count", 0)
        returned_count = result.get("returned_count", 0)
        
        if not assets:
            return f"No assets found for type '{asset_type}'" + (f" in categories '{categories}'" if categories else "")
        
        formatted_output = f"Found {total_count} assets"
        if categories:
            formatted_output += f" in categories: {categories}"
        formatted_output += f"\\nShowing {returned_count} assets:\\n\\n"
        
        # Sort assets by download count (popularity)
        sorted_assets = sorted(assets.items(), key=lambda x: x[1].get("download_count", 0), reverse=True)
        
        for asset_id, asset_data in sorted_assets:
            formatted_output += f"- {asset_data.get('name', asset_id)} (ID: {asset_id})\\n"
            formatted_output += f"  Type: {['HDRI', 'Texture', 'Model'][asset_data.get('type', 0)]}\\n"
            formatted_output += f"  Categories: {', '.join(asset_data.get('categories', []))}\\n"
            formatted_output += f"  Downloads: {asset_data.get('download_count', 'Unknown')}\\n\\n"
        
        return formatted_output


class PolyHavenDownloadTool(BaseTool):
    """Download PolyHaven assets."""
    
    @tool_wrapper
    def execute(self, asset_id: str, asset_type: str, resolution: str = "1k", file_format: Optional[str] = None) -> str:
        """Download and import a PolyHaven asset into Blender."""
        # Validate parameters
        validated_params = validate_asset_params(asset_id, asset_type, resolution)
        
        result = self.send_command("download_polyhaven_asset", {
            "asset_id": validated_params['asset_id'],
            "asset_type": validated_params['asset_type'],
            "resolution": validated_params['resolution'],
            "file_format": file_format
        })
        
        if "error" in result:
            return f"Error: {result['error']}"
        
        if result.get("success"):
            message = result.get("message", "Asset downloaded and imported successfully")
            
            # Add additional information based on asset type
            if asset_type == "hdris":
                return f"{message}. The HDRI has been set as the world environment."
            elif asset_type == "textures":
                material_name = result.get("material", "")
                maps = ", ".join(result.get("maps", []))
                return f"{message}. Created material '{material_name}' with maps: {maps}."
            elif asset_type == "models":
                return f"{message}. The model has been imported into the current scene."
            else:
                return message
        else:
            return f"Failed to download asset: {result.get('message', 'Unknown error')}"


class PolyHavenTextureTool(BaseTool):
    """Apply PolyHaven textures to objects."""
    
    @tool_wrapper
    def execute(self, object_name: str, texture_id: str) -> str:
        """Apply a previously downloaded PolyHaven texture to an object."""
        # Validate inputs
        validated_object_name = validate_object_name(object_name)
        validated_texture_id = validate_asset_params(texture_id, "textures", "1k")['asset_id']
        
        result = self.send_command("set_texture", {
            "object_name": validated_object_name,
            "texture_id": validated_texture_id
        })
        
        if "error" in result:
            return f"Error: {result['error']}"
        
        if result.get("success"):
            material_name = result.get("material", "")
            maps = ", ".join(result.get("maps", []))
            
            # Add detailed material info
            material_info = result.get("material_info", {})
            node_count = material_info.get("node_count", 0)
            has_nodes = material_info.get("has_nodes", False)
            texture_nodes = material_info.get("texture_nodes", [])
            
            output = f"Successfully applied texture '{validated_texture_id}' to {validated_object_name}.\\n"
            output += f"Using material '{material_name}' with maps: {maps}.\\n\\n"
            output += f"Material has nodes: {has_nodes}\\n"
            output += f"Total node count: {node_count}\\n\\n"
            
            if texture_nodes:
                output += "Texture nodes:\\n"
                for node in texture_nodes:
                    output += f"- {node['name']} using image: {node['image']}\\n"
                    if node.get('connections'):
                        output += "  Connections:\\n"
                        for conn in node['connections']:
                            output += f"    {conn}\\n"
            else:
                output += "No texture nodes found in the material.\\n"
            
            return output
        else:
            return f"Failed to apply texture: {result.get('message', 'Unknown error')}"


# Tool instances
polyhaven_status_tool = PolyHavenStatusTool()
polyhaven_categories_tool = PolyHavenCategoriesTool()
polyhaven_search_tool = PolyHavenSearchTool()
polyhaven_download_tool = PolyHavenDownloadTool()
polyhaven_texture_tool = PolyHavenTextureTool()


# MCP tool function wrappers
def get_polyhaven_status(ctx: Context) -> str:
    """Check if PolyHaven integration is enabled in Blender."""
    return polyhaven_status_tool.execute()


def get_polyhaven_categories(ctx: Context, asset_type: str = "hdris") -> str:
    """Get a list of categories for a specific asset type on PolyHaven."""
    return polyhaven_categories_tool.execute(asset_type=asset_type)


def search_polyhaven_assets(ctx: Context, asset_type: str = "all", categories: Optional[str] = None) -> str:
    """Search for assets on PolyHaven with optional filtering."""
    return polyhaven_search_tool.execute(asset_type=asset_type, categories=categories)


def download_polyhaven_asset(ctx: Context, asset_id: str, asset_type: str, resolution: str = "1k", file_format: Optional[str] = None) -> str:
    """Download and import a PolyHaven asset into Blender."""
    return polyhaven_download_tool.execute(asset_id=asset_id, asset_type=asset_type, resolution=resolution, file_format=file_format)


def set_texture(ctx: Context, object_name: str, texture_id: str) -> str:
    """Apply a previously downloaded PolyHaven texture to an object."""
    return polyhaven_texture_tool.execute(object_name=object_name, texture_id=texture_id)
