# 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)

"""
Hyper3D Rodin integration tools for the Blender MCP server.

This module provides tools for generating 3D models using Hyper3D Rodin AI technology
and importing them into Blender. It supports both text-based and image-based generation,
job status polling, and asset importing with comprehensive error handling.

Features:
- Hyper3D Rodin service status verification
- Text-to-3D model generation with customizable parameters
- Image-to-3D model generation with multiple input modes
- Job status polling for generation progress tracking
- Generated asset importing with proper scene integration
- Comprehensive validation and security measures
"""

import json
import os
import base64
from pathlib import Path
from urllib.parse import urlparse
from typing import Optional, List, Union, Dict, Any

from ..connection import get_blender_connection
from ..utils import get_logger, validate_string_input, sanitize_text_input
from ..exceptions import BlenderMCPError, ConnectionError, ValidationError
from .base import blender_tool

# Configure logging
logger = get_logger(__name__)


@blender_tool()
def get_hyper3d_status() -> str:
    """
    Check if Hyper3D Rodin integration is enabled in Blender.
    
    Verifies whether the Hyper3D Rodin addon is properly installed and configured
    in Blender, enabling AI-powered 3D model generation capabilities.
    
    Returns:
        str: Status message indicating whether Hyper3D Rodin features are available,
             including configuration details and mode information.
    
    Raises:
        ConnectionError: If unable to connect to Blender
        BlenderMCPError: If there's an error checking Hyper3D status
    
    Note:
        The returned message contains sensitive API configuration details
        that should be handled securely and not exposed to unauthorized users.
    """
    try:
        logger.info("Checking Hyper3D Rodin integration status")
        blender = get_blender_connection()
        
        result = blender.send_command("get_hyper3d_status")
        
        if not isinstance(result, dict):
            raise BlenderMCPError("Invalid response from Blender for Hyper3D status")
        
        enabled = result.get("enabled", False)
        message = result.get("message", "")
        
        logger.info(f"Hyper3D status check completed: enabled={enabled}")
        return message
        
    except ConnectionError:
        raise
    except Exception as e:
        error_msg = f"Error checking Hyper3D status: {str(e)}"
        logger.error(error_msg)
        raise BlenderMCPError(error_msg) from e


@blender_tool()
def generate_hyper3d_model_via_text(
    text_prompt: str,
    bbox_condition: Optional[List[float]] = None
) -> str:
    """
    Generate a 3D model using Hyper3D Rodin from a text description.
    
    Creates a 3D asset using AI-powered generation based on a textual description.
    The generated model includes built-in materials and is normalized in size,
    so rescaling after generation may be useful for scene integration.
    
    Args:
        text_prompt: A concise description of the desired model in English.
                    Should be specific and descriptive for best results.
        bbox_condition: Optional list of 3 floats controlling the ratio between
                       [Length, Width, Height] of the generated model.
                       Values should be positive and represent relative proportions.
    
    Returns:
        str: JSON string containing task identifiers for status polling:
             - task_uuid: Unique identifier for the generation task
             - subscription_key: Key for polling job status
             Or error details if generation initiation fails.
    
    Raises:
        ValidationError: If input parameters are invalid
        ConnectionError: If unable to connect to Blender
        BlenderMCPError: If generation task creation fails
    
    Security Notes:
        - Text prompt is sanitized to prevent injection attacks
        - Bbox parameters are validated for safe numeric ranges
        - API keys and sensitive data are handled securely
    
    Example:
        >>> generate_hyper3d_model_via_text("red sports car", [2.0, 1.0, 0.8])
        {"task_uuid": "abc123", "subscription_key": "key456"}
    """
    try:
        # Input validation and sanitization
        text_prompt = validate_string_input(text_prompt, "text_prompt", min_length=3, max_length=500)
        text_prompt = sanitize_text_input(text_prompt)
        
        # Validate and process bbox_condition
        processed_bbox = _process_bbox_condition(bbox_condition)
        
        logger.info(f"Initiating Hyper3D text-to-3D generation: prompt='{text_prompt[:50]}...'")
        
        blender = get_blender_connection()
        
        # Send generation command to Blender
        result = blender.send_command("create_rodin_job", {
            "text_prompt": text_prompt,
            "images": None,
            "bbox_condition": processed_bbox,
        })
        
        if not isinstance(result, dict):
            raise BlenderMCPError("Invalid response from Blender for Hyper3D generation")
        
        # Check if job submission was successful
        submit_time = result.get("submit_time")
        if submit_time:
            response_data = {
                "task_uuid": result["uuid"],
                "subscription_key": result["jobs"]["subscription_key"],
            }
            logger.info(f"Hyper3D generation task created successfully: {result['uuid']}")
            return json.dumps(response_data)
        else:
            logger.warning(f"Hyper3D generation task creation failed: {result}")
            return json.dumps(result)
            
    except (ValidationError, ConnectionError):
        raise
    except Exception as e:
        error_msg = f"Error generating Hyper3D task: {str(e)}"
        logger.error(error_msg, exc_info=True)
        raise BlenderMCPError(error_msg) from e


@blender_tool()
def generate_hyper3d_model_via_images(
    input_image_paths: Optional[List[str]] = None,
    input_image_urls: Optional[List[str]] = None,
    bbox_condition: Optional[List[float]] = None
) -> str:
    """
    Generate a 3D model using Hyper3D Rodin from input images.
    
    Creates a 3D asset using AI-powered generation based on one or more input images.
    Supports both local file paths and remote URLs depending on the Hyper3D Rodin mode.
    The generated model includes built-in materials and normalized sizing.
    
    Args:
        input_image_paths: Absolute paths to local input images. Required for MAIN_SITE mode.
                          Even single images should be provided as a list.
        input_image_urls: URLs of remote input images. Required for FAL_AI mode.
                         Even single images should be provided as a list.
        bbox_condition: Optional list of 3 floats controlling the ratio between
                       [Length, Width, Height] of the generated model.
    
    Returns:
        str: JSON string containing task identifiers for status polling or error details
    
    Raises:
        ValidationError: If input parameters are invalid or conflicting
        ConnectionError: If unable to connect to Blender
        BlenderMCPError: If generation task creation fails
    
    Security Notes:
        - File paths are validated for existence and safety
        - URLs are validated for proper format
        - Image data is encoded securely using base64
        - File access is restricted to specified paths only
    
    Usage:
        Only one of input_image_paths or input_image_urls should be provided,
        depending on the current Hyper3D Rodin mode configuration.
    
    Example:
        >>> generate_hyper3d_model_via_images(
        ...     input_image_paths=["/path/to/image1.jpg", "/path/to/image2.png"],
        ...     bbox_condition=[1.0, 1.0, 1.5]
        ... )
        {"task_uuid": "def789", "subscription_key": "key123"}
    """
    try:
        # Validate input parameter combinations
        if input_image_paths is not None and input_image_urls is not None:
            raise ValidationError("Cannot specify both input_image_paths and input_image_urls")
        
        if input_image_paths is None and input_image_urls is None:
            raise ValidationError("Must specify either input_image_paths or input_image_urls")
        
        # Process and validate image inputs
        if input_image_paths is not None:
            images = _process_image_paths(input_image_paths)
            logger.info(f"Processing {len(input_image_paths)} local image files")
        else:
            images = _process_image_urls(input_image_urls)
            logger.info(f"Processing {len(input_image_urls)} remote image URLs")
        
        # Validate and process bbox_condition
        processed_bbox = _process_bbox_condition(bbox_condition)
        
        logger.info("Initiating Hyper3D image-to-3D generation")
        
        blender = get_blender_connection()
        
        # Send generation command to Blender
        result = blender.send_command("create_rodin_job", {
            "text_prompt": None,
            "images": images,
            "bbox_condition": processed_bbox,
        })
        
        if not isinstance(result, dict):
            raise BlenderMCPError("Invalid response from Blender for Hyper3D generation")
        
        # Check if job submission was successful
        submit_time = result.get("submit_time")
        if submit_time:
            response_data = {
                "task_uuid": result["uuid"],
                "subscription_key": result["jobs"]["subscription_key"],
            }
            logger.info(f"Hyper3D generation task created successfully: {result['uuid']}")
            return json.dumps(response_data)
        else:
            logger.warning(f"Hyper3D generation task creation failed: {result}")
            return json.dumps(result)
            
    except (ValidationError, ConnectionError):
        raise
    except Exception as e:
        error_msg = f"Error generating Hyper3D task: {str(e)}"
        logger.error(error_msg, exc_info=True)
        raise BlenderMCPError(error_msg) from e


@blender_tool()
def poll_rodin_job_status(
    subscription_key: Optional[str] = None,
    request_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Check the status of a Hyper3D Rodin generation task.
    
    Polls the generation service to check if a 3D model generation task has completed.
    Supports both MAIN_SITE and FAL_AI modes with different parameter requirements.
    This is a polling API that should be called periodically until completion.
    
    Args:
        subscription_key: The subscription key from the generation task (MAIN_SITE mode)
        request_id: The request ID from the generation task (FAL_AI mode)
    
    Returns:
        Dict[str, Any]: Status information with different formats based on mode:
        
        MAIN_SITE mode:
        - List of status strings. Task is complete when all are "Done"
        - "Failed" indicates generation failure
        - "Canceled" indicates task cancellation
        
        FAL_AI mode:
        - Status field with values: "COMPLETED", "IN_PROGRESS", "IN_QUEUE"
        - Additional progress and error information if available
    
    Raises:
        ValidationError: If neither or both parameters are provided
        ConnectionError: If unable to connect to Blender
        BlenderMCPError: If status polling fails
    
    Usage Notes:
        - Only provide one of subscription_key or request_id based on the Hyper3D mode
        - Continue polling until status indicates completion or failure
        - Avoid excessive polling frequency to prevent service overload
    
    Example:
        >>> poll_rodin_job_status(subscription_key="key123")
        {"status": ["Processing", "Done", "Done"]}
    """
    try:
        # Validate input parameters
        if subscription_key is not None and request_id is not None:
            raise ValidationError("Cannot specify both subscription_key and request_id")
        
        if subscription_key is None and request_id is None:
            raise ValidationError("Must specify either subscription_key or request_id")
        
        # Prepare parameters based on mode
        if subscription_key is not None:
            subscription_key = validate_string_input(subscription_key, "subscription_key", min_length=1, max_length=200)
            kwargs = {"subscription_key": subscription_key}
            logger.info(f"Polling Hyper3D job status (MAIN_SITE mode): {subscription_key[:20]}...")
        else:
            request_id = validate_string_input(request_id, "request_id", min_length=1, max_length=200)
            kwargs = {"request_id": request_id}
            logger.info(f"Polling Hyper3D job status (FAL_AI mode): {request_id[:20]}...")
        
        blender = get_blender_connection()
        
        # Send status polling command to Blender
        result = blender.send_command("poll_rodin_job_status", kwargs)
        
        if not isinstance(result, dict):
            # Handle case where result might be a simple status string
            if isinstance(result, str):
                return {"status": result}
            raise BlenderMCPError("Invalid response from Blender for Hyper3D status polling")
        
        logger.debug(f"Hyper3D job status: {result}")
        return result
        
    except (ValidationError, ConnectionError):
        raise
    except Exception as e:
        error_msg = f"Error polling Hyper3D job status: {str(e)}"
        logger.error(error_msg, exc_info=True)
        raise BlenderMCPError(error_msg) from e


@blender_tool()
def import_generated_asset(
    name: str,
    task_uuid: Optional[str] = None,
    request_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Import a generated 3D asset from Hyper3D Rodin into the Blender scene.
    
    Retrieves and imports a completed 3D model from the Hyper3D Rodin service
    into the current Blender scene. The asset is imported with the specified name
    and includes all materials and textures generated by the AI.
    
    Args:
        name: The name to assign to the imported object in the Blender scene.
              Should be a valid Blender object name.
        task_uuid: Unique task identifier for MAIN_SITE mode
        request_id: Request identifier for FAL_AI mode
    
    Returns:
        Dict[str, Any]: Import result information including:
        - success: Boolean indicating import success
        - imported_objects: List of created object names
        - message: Status or error message
        - Additional metadata about the imported asset
    
    Raises:
        ValidationError: If parameters are invalid or conflicting
        ConnectionError: If unable to connect to Blender
        BlenderMCPError: If asset import fails
    
    Security Notes:
        - Asset data is validated before import
        - Object names are sanitized to prevent conflicts
        - Import process includes safety checks for malicious content
    
    Usage:
        Only provide one of task_uuid or request_id based on the Hyper3D mode
        used for generation. Ensure the generation task is completed before
        attempting to import the asset.
    
    Example:
        >>> import_generated_asset("MyGeneratedModel", task_uuid="abc123")
        {"success": True, "imported_objects": ["MyGeneratedModel"], "message": "Import successful"}
    """
    try:
        # Validate input parameters
        name = validate_string_input(name, "name", min_length=1, max_length=100)
        name = sanitize_text_input(name, allow_special_chars=True)
        
        if task_uuid is not None and request_id is not None:
            raise ValidationError("Cannot specify both task_uuid and request_id")
        
        if task_uuid is None and request_id is None:
            raise ValidationError("Must specify either task_uuid or request_id")
        
        # Prepare parameters based on mode
        kwargs = {"name": name}
        
        if task_uuid is not None:
            task_uuid = validate_string_input(task_uuid, "task_uuid", min_length=1, max_length=200)
            kwargs["task_uuid"] = task_uuid
            logger.info(f"Importing generated asset (MAIN_SITE mode): {name}")
        else:
            request_id = validate_string_input(request_id, "request_id", min_length=1, max_length=200)
            kwargs["request_id"] = request_id
            logger.info(f"Importing generated asset (FAL_AI mode): {name}")
        
        blender = get_blender_connection()
        
        # Send import command to Blender
        result = blender.send_command("import_generated_asset", kwargs)
        
        if not isinstance(result, dict):
            raise BlenderMCPError("Invalid response from Blender for asset import")
        
        # Log import result
        if result.get("success"):
            imported_objects = result.get("imported_objects", [])
            logger.info(f"Asset import successful: {len(imported_objects)} objects created")
        else:
            logger.warning(f"Asset import failed: {result.get('message', 'Unknown error')}")
        
        return result
        
    except (ValidationError, ConnectionError):
        raise
    except Exception as e:
        error_msg = f"Error importing generated asset: {str(e)}"
        logger.error(error_msg, exc_info=True)
        raise BlenderMCPError(error_msg) from e


def _process_bbox_condition(bbox_condition: Optional[List[float]]) -> Optional[List[int]]:
    """
    Process and validate bbox_condition parameter for Hyper3D generation.
    
    Args:
        bbox_condition: Optional list of 3 floats representing [Length, Width, Height] ratios
    
    Returns:
        Optional[List[int]]: Processed bbox condition as integers, or None if input is None
    
    Raises:
        ValidationError: If bbox_condition format or values are invalid
    """
    if bbox_condition is None:
        return None
    
    # Validate input type and length
    if not isinstance(bbox_condition, list) or len(bbox_condition) != 3:
        raise ValidationError("bbox_condition must be a list of exactly 3 numbers")
    
    # Validate all values are numeric and positive
    try:
        float_values = [float(x) for x in bbox_condition]
    except (ValueError, TypeError) as e:
        raise ValidationError(f"bbox_condition values must be numeric: {str(e)}") from e
    
    # Check for positive values
    if any(val <= 0 for val in float_values):
        raise ValidationError("bbox_condition values must be greater than zero")
    
    # Check for reasonable bounds (prevent extreme ratios)
    if any(val > 1000 for val in float_values):
        raise ValidationError("bbox_condition values are too large (max: 1000)")
    
    # Convert to normalized integer ratios
    max_val = max(float_values)
    normalized = [int(float(val) / max_val * 100) for val in float_values]
    
    logger.debug(f"Processed bbox_condition: {bbox_condition} -> {normalized}")
    return normalized


def _process_image_paths(image_paths: List[str]) -> List[tuple]:
    """
    Process and validate local image file paths for Hyper3D generation.
    
    Args:
        image_paths: List of absolute paths to image files
    
    Returns:
        List[tuple]: List of (file_extension, base64_encoded_data) tuples
    
    Raises:
        ValidationError: If any image path is invalid or inaccessible
    """
    if not isinstance(image_paths, list) or not image_paths:
        raise ValidationError("input_image_paths must be a non-empty list")
    
    if len(image_paths) > 10:  # Reasonable limit
        raise ValidationError("Too many image files (max: 10)")
    
    images = []
    allowed_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    
    for path in image_paths:
        # Validate path format
        path = validate_string_input(path, "image_path", min_length=1, max_length=500)
        
        # Check file existence and accessibility
        if not os.path.exists(path):
            raise ValidationError(f"Image file does not exist: {path}")
        
        if not os.path.isfile(path):
            raise ValidationError(f"Path is not a file: {path}")
        
        # Validate file extension
        path_obj = Path(path)
        if path_obj.suffix.lower() not in allowed_extensions:
            raise ValidationError(f"Unsupported image format: {path_obj.suffix}")
        
        # Check file size (prevent extremely large files)
        file_size = os.path.getsize(path)
        if file_size > 50 * 1024 * 1024:  # 50MB limit
            raise ValidationError(f"Image file too large: {path} ({file_size} bytes)")
        
        # Read and encode image data
        try:
            with open(path, "rb") as f:
                image_data = f.read()
            encoded_data = base64.b64encode(image_data).decode("ascii")
            images.append((path_obj.suffix, encoded_data))
            logger.debug(f"Encoded image: {path} ({len(encoded_data)} chars)")
        except (IOError, OSError) as e:
            raise ValidationError(f"Error reading image file {path}: {str(e)}") from e
        except Exception as e:
            raise ValidationError(f"Error processing image file {path}: {str(e)}") from e
    
    return images


def _process_image_urls(image_urls: List[str]) -> List[str]:
    """
    Process and validate remote image URLs for Hyper3D generation.
    
    Args:
        image_urls: List of URLs pointing to image files
    
    Returns:
        List[str]: List of validated and sanitized URLs
    
    Raises:
        ValidationError: If any URL is invalid or unsafe
    """
    if not isinstance(image_urls, list) or not image_urls:
        raise ValidationError("input_image_urls must be a non-empty list")
    
    if len(image_urls) > 10:  # Reasonable limit
        raise ValidationError("Too many image URLs (max: 10)")
    
    validated_urls = []
    allowed_schemes = {'http', 'https'}
    allowed_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
    
    for url in image_urls:
        # Validate URL format
        url = validate_string_input(url, "image_url", min_length=1, max_length=2000)
        url = sanitize_text_input(url, allow_special_chars=True)
        
        # Parse and validate URL components
        try:
            parsed = urlparse(url)
        except Exception as e:
            raise ValidationError(f"Invalid URL format: {url}") from e
        
        if not parsed.scheme or parsed.scheme not in allowed_schemes:
            raise ValidationError(f"Invalid URL scheme (must be http/https): {url}")
        
        if not parsed.netloc:
            raise ValidationError(f"Invalid URL (missing domain): {url}")
        
        # Check for reasonable URL length and structure
        if len(url) > 2000:
            raise ValidationError(f"URL too long: {url}")
        
        # Validate file extension if present in path
        if parsed.path:
            path_obj = Path(parsed.path)
            if path_obj.suffix and path_obj.suffix.lower() not in allowed_extensions:
                logger.warning(f"URL may not point to supported image format: {url}")
        
        validated_urls.append(url)
        logger.debug(f"Validated image URL: {url}")
    
    return validated_urls
