"""
MCP Tools implementation using direct library imports.

This module implements MCP tools by calling coaiapy, langfuse, and redis libraries directly
instead of using subprocess wrappers. This provides:
- Faster execution (no process creation overhead)
- Better error handling (typed exceptions)
- Shared configuration (load once, use everywhere)
- No environment variable propagation issues
"""

import os
from typing import Dict, Any, Optional, List
import redis
from langfuse import Langfuse

# Import from coaiapy
try:
    from coaiapy import coaiamodule
    from coaiapy.cofuse import (
        list_score_configs,
        get_score_config,
        apply_score_config,
        create_score_for_target,
        list_prompts as cofuse_list_prompts,
        get_prompt as cofuse_get_prompt,
        create_prompt as cofuse_create_prompt,
        update_prompt_version_labels as cofuse_update_prompt_version_labels,
        list_datasets as cofuse_list_datasets,
        get_dataset as cofuse_get_dataset,
        add_trace,
        add_observation,
        patch_trace_output,
        get_observation,
        format_observation_display,
        get_comments,
        get_comment_by_id,
        post_comment,
        list_traces,
        list_projects,
        get_trace_with_observations,
        format_traces_table,
        format_trace_tree,
        upload_and_attach_media,
        get_media,
        format_media_display,
    )
    from coaiapy.pipeline import TemplateLoader
except ImportError as e:
    print(f"Warning: Could not import from coaiapy: {e}")
    print("Some tools may not be available.")

# Load configuration once on module import
# Support custom .env path via COAIAPY_ENV_PATH environment variable
try:
    env_path = os.getenv('COAIAPY_ENV_PATH')
    config = coaiamodule.read_config(env_path=env_path)
except Exception as e:
    print(f"Warning: Could not load config: {e}")
    config = {}

# Initialize Redis client
# Try direct URL first (handles SSL automatically), fall back to component config
redis_client = None
REDIS_AVAILABLE = False

# Check for direct Redis URL from environment
redis_url = os.getenv('REDIS_URL') or os.getenv('KV_URL')
if redis_url:
    try:
        redis_client = redis.from_url(redis_url, decode_responses=True)
        redis_client.ping()
        REDIS_AVAILABLE = True
    except Exception as e:
        print(f"Warning: Redis connection failed with URL: {e}")
        redis_client = None

# Fall back to component-based config if URL didn't work
if not REDIS_AVAILABLE:
    redis_config = config.get("jtaleconf", {})
    try:
        redis_client = redis.Redis(
            host=redis_config.get("host", "localhost"),
            port=redis_config.get("port", 6379),
            db=redis_config.get("db", 0),
            password=redis_config.get("password") if redis_config.get("password") else None,
            ssl=redis_config.get("ssl", False),
            ssl_cert_reqs="none" if redis_config.get("ssl") else "required",
            decode_responses=True,
        )
        # Test connection
        redis_client.ping()
        REDIS_AVAILABLE = True
    except (redis.RedisError, redis.ConnectionError) as e:
        print(f"Warning: Redis not available: {e}")
        redis_client = None
        REDIS_AVAILABLE = False

# Initialize Langfuse client
try:
    langfuse_client = Langfuse(
        secret_key=config.get("langfuse_secret_key", os.getenv("LANGFUSE_SECRET_KEY")),
        public_key=config.get("langfuse_public_key", os.getenv("LANGFUSE_PUBLIC_KEY")),
        host=config.get("langfuse_host", os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com")),
    )
    LANGFUSE_AVAILABLE = True
except Exception as e:
    print(f"Warning: Langfuse not available: {e}")
    langfuse_client = None
    LANGFUSE_AVAILABLE = False

# Initialize Pipeline Template Loader
try:
    pipeline_loader = TemplateLoader()
    PIPELINE_AVAILABLE = True
except Exception as e:
    print(f"Warning: Pipeline loader not available: {e}")
    pipeline_loader = None
    PIPELINE_AVAILABLE = False


# ============================================================================
# Redis Tools
# ============================================================================

async def coaia_tash(key: str, value: str) -> Dict[str, Any]:
    """
    Stash key-value pair to Redis via direct client call.
    
    Args:
        key: Redis key
        value: Value to store
        
    Returns:
        Dict with success status and message/error
    """
    if not REDIS_AVAILABLE:
        return {
            "success": False,
            "error": "Redis is not available. Check configuration and Redis server."
        }
    
    try:
        redis_client.set(key, value)
        return {
            "success": True,
            "message": f"Stored '{key}' in Redis"
        }
    except redis.RedisError as e:
        return {
            "success": False,
            "error": f"Redis error: {str(e)}"
        }


async def coaia_fetch(key: str) -> Dict[str, Any]:
    """
    Fetch value from Redis via direct client call.
    
    Args:
        key: Redis key to fetch
        
    Returns:
        Dict with success status and value/error
    """
    if not REDIS_AVAILABLE:
        return {
            "success": False,
            "error": "Redis is not available. Check configuration and Redis server."
        }
    
    try:
        value = redis_client.get(key)
        if value is None:
            return {
                "success": False,
                "error": f"Key '{key}' not found in Redis"
            }
        return {
            "success": True,
            "value": value
        }
    except redis.RedisError as e:
        return {
            "success": False,
            "error": f"Redis error: {str(e)}"
        }


# ============================================================================
# Langfuse Trace Tools
# ============================================================================

async def coaia_fuse_trace_create(
    trace_id: str,
    user_id: Optional[str] = None,
    session_id: Optional[str] = None,
    name: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    input_data: Optional[Any] = None,
    output_data: Optional[Any] = None,
) -> Dict[str, Any]:
    """
    Create Langfuse trace via direct SDK call.
    
    Args:
        trace_id: Unique trace identifier
        user_id: Optional user identifier
        session_id: Optional session identifier
        name: Optional trace name
        metadata: Optional metadata dictionary
        input_data: Optional input data
        output_data: Optional output data
        
    Returns:
        Dict with success status and trace details/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's add_trace function which handles the API call
        result = add_trace(
            trace_id=trace_id,
            user_id=user_id,
            session_id=session_id,
            name=name,
            input_data=input_data,
            output_data=output_data,
            metadata=metadata,
        )
        
        return {
            "success": True,
            "trace_id": trace_id,
            "details": {
                "name": name,
                "user_id": user_id,
                "session_id": session_id,
                "metadata": metadata,
            }
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse trace creation error: {str(e)}"
        }


async def coaia_fuse_add_observation(
    observation_id: str,
    trace_id: str,
    name: str,
    observation_type: str = "SPAN",
    parent_id: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    input_data: Optional[Any] = None,
    output_data: Optional[Any] = None,
    start_time: Optional[str] = None,
    end_time: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Add observation to trace via direct SDK call.
    
    Args:
        observation_id: Unique observation identifier
        trace_id: Parent trace identifier
        name: Observation name
        observation_type: Type (SPAN, EVENT, GENERATION)
        parent_id: Optional parent observation ID
        metadata: Optional metadata dictionary
        input_data: Optional input data
        output_data: Optional output data
        start_time: Optional start timestamp
        end_time: Optional end timestamp
        
    Returns:
        Dict with success status and observation details/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's add_observation function
        result = add_observation(
            observation_id=observation_id,
            trace_id=trace_id,
            observation_type=observation_type,
            name=name,
            parent_observation_id=parent_id,
            metadata=metadata,
            input_data=input_data,
            output_data=output_data,
            start_time=start_time,
            end_time=end_time,
        )
        
        return {
            "success": True,
            "observation_id": observation_id,
            "trace_id": trace_id,
            "name": name,
            "type": observation_type,
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse observation creation error: {str(e)}"
        }


async def coaia_fuse_trace_patch_output(
    trace_id: str,
    output_data: Any,
) -> Dict[str, Any]:
    """
    Update the output field of an existing trace.

    Args:
        trace_id: Trace identifier to update
        output_data: New output data (string, object, or any JSON-serializable data)

    Returns:
        Dict with success status and trace details/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }

    try:
        # Use coaiapy's patch_trace_output function
        result = patch_trace_output(
            trace_id=trace_id,
            output_data=output_data,
        )

        return {
            "success": True,
            "trace_id": trace_id,
            "message": f"Successfully patched output for trace {trace_id}",
            "details": {
                "output_data": output_data if not isinstance(output_data, str) or len(str(output_data)) < 100 else f"{str(output_data)[:100]}...",
            }
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse trace output patch error: {str(e)}"
        }


async def coaia_fuse_trace_get(trace_id: str, json_output: bool = False) -> Dict[str, Any]:
    """
    Get a specific trace by ID from Langfuse with all its observations.

    Args:
        trace_id: Unique trace identifier
        json_output: If True, return raw JSON; if False, return formatted tree

    Returns:
        Dict with success status and trace data/error with proper Langfuse URL
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }

    try:
        # Fetch trace with observations
        trace_data = get_trace_with_observations(trace_id)

        import json
        parsed = json.loads(trace_data)
        if 'error' in parsed:
            return {
                "success": False,
                "error": parsed['error']
            }

        # Get project ID for constructing proper URL
        try:
            projects_json = list_projects()
            projects = json.loads(projects_json)
            # Projects response has structure: {"data": [{"id": "...", "name": "...", ...}]}
            project_id = None
            if isinstance(projects, dict) and 'data' in projects and isinstance(projects['data'], list) and len(projects['data']) > 0:
                project_id = projects['data'][0].get('id')

            # Construct proper Langfuse URL with project_id
            langfuse_host = config.get('langfuse_host', 'https://cloud.langfuse.com')
            trace_url = f"{langfuse_host}/project/{project_id}/traces/{trace_id}" if project_id else f"{langfuse_host}/traces/{trace_id}"
        except Exception as url_error:
            # Fallback if project fetch fails
            langfuse_host = config.get('langfuse_host', 'https://cloud.langfuse.com')
            trace_url = f"{langfuse_host}/traces/{trace_id}"

        if json_output:
            return {
                "success": True,
                "trace": parsed,
                "trace_url": trace_url,
                "json": trace_data
            }
        else:
            formatted = format_trace_tree(parsed)
            return {
                "success": True,
                "trace": parsed,
                "trace_url": trace_url,
                "formatted": formatted
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Trace retrieval error: {str(e)}"
        }


async def coaia_fuse_trace_view(trace_id: str, json_output: bool = False) -> Dict[str, Any]:
    """
    View trace details with observations from Langfuse (alias for coaia_fuse_trace_get).

    Args:
        trace_id: Trace identifier to fetch
        json_output: If True, return raw JSON; if False, return formatted tree

    Returns:
        Dict with success status and trace data/error
    """
    return await coaia_fuse_trace_get(trace_id, json_output)


async def coaia_fuse_observation_get(observation_id: str, json_output: bool = False) -> Dict[str, Any]:
    """
    Get a specific observation by ID from Langfuse.

    Args:
        observation_id: Unique observation identifier
        json_output: If True, return raw JSON; if False, return formatted display

    Returns:
        Dict with success status and observation data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }

    try:
        # Use coaiapy's get_observation function
        obs_data = get_observation(observation_id)

        import json
        parsed = json.loads(obs_data)
        if 'error' in parsed:
            return {
                "success": False,
                "error": parsed['error']
            }

        if json_output:
            return {
                "success": True,
                "observation": parsed,
                "json": obs_data
            }
        else:
            formatted = format_observation_display(parsed)
            return {
                "success": True,
                "observation": parsed,
                "formatted": formatted
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Observation retrieval error: {str(e)}"
        }


async def coaia_fuse_traces_list(
    session_id: Optional[str] = None,
    user_id: Optional[str] = None,
    name: Optional[str] = None,
    tags: Optional[List[str]] = None,
    from_timestamp: Optional[str] = None,
    to_timestamp: Optional[str] = None,
    order_by: Optional[str] = None,
    version: Optional[str] = None,
    release: Optional[str] = None,
    environment: Optional[List[str]] = None,
    page: Optional[int] = 1,
    limit: Optional[int] = 50,
    json_output: bool = False
) -> Dict[str, Any]:
    """
    List traces with comprehensive filtering options.
    
    Args:
        session_id: Filter by session ID
        user_id: Filter by user ID
        name: Filter by trace name (exact match)
        tags: List of tags - only traces with ALL tags will be returned
        from_timestamp: Include traces from this timestamp (ISO 8601 format)
        to_timestamp: Include traces before this timestamp (ISO 8601 format)
        order_by: Sort order, format: field.direction (e.g., "timestamp.asc", "timestamp.desc")
                 Fields: id, timestamp, name, userId, release, version, sessionId
        version: Filter by version
        release: Filter by release
        environment: List of environment values
        page: Page number (starts at 1)
        limit: Items per page (default 50)
        json_output: Return raw JSON instead of formatted table
        
    Returns:
        Dict with success status and traces data or formatted table
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's list_traces function with all filters
        traces_data = list_traces(
            include_observations=False,
            session_id=session_id,
            user_id=user_id,
            name=name,
            tags=tags,
            from_timestamp=from_timestamp,
            to_timestamp=to_timestamp,
            order_by=order_by,
            version=version,
            release=release,
            environment=environment,
            page=page,
            limit=limit
        )
        
        import json
        parsed = json.loads(traces_data)
        
        if json_output:
            return {
                "success": True,
                "traces": parsed,
                "filters": {
                    "session_id": session_id,
                    "user_id": user_id,
                    "name": name,
                    "tags": tags,
                    "from_timestamp": from_timestamp,
                    "to_timestamp": to_timestamp,
                    "order_by": order_by,
                    "version": version,
                    "release": release,
                    "environment": environment,
                    "page": page,
                    "limit": limit
                }
            }
        else:
            # Format as table
            formatted = format_traces_table(parsed)
            return {
                "success": True,
                "formatted": formatted,
                "traces": parsed,
                "filters": {
                    "session_id": session_id,
                    "user_id": user_id,
                    "name": name,
                    "tags": tags,
                    "from_timestamp": from_timestamp,
                    "to_timestamp": to_timestamp,
                    "order_by": order_by,
                    "version": version,
                    "release": release,
                    "environment": environment,
                    "page": page,
                    "limit": limit
                }
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Trace listing error: {str(e)}"
        }


async def coaia_fuse_traces_session_view(session_id: str, json_output: bool = False) -> Dict[str, Any]:
    """
    View all traces for a specific session from Langfuse.

    Args:
        session_id: Session identifier to filter traces
        json_output: If True, return raw JSON data; if False, return formatted table

    Returns:
        Dict with success status and traces data or formatted table/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }

    try:
        import json

        # Get project ID for constructing proper URL
        try:
            projects_json = list_projects()
            projects = json.loads(projects_json)
            # Projects response has structure: {"data": [{"id": "...", "name": "...", ...}]}
            project_id = None
            if isinstance(projects, dict) and 'data' in projects and isinstance(projects['data'], list) and len(projects['data']) > 0:
                project_id = projects['data'][0].get('id')

            # Construct proper Langfuse URL with project_id
            langfuse_host = config.get('langfuse_host', 'https://cloud.langfuse.com')
            session_url = f"{langfuse_host}/project/{project_id}/sessions/{session_id}" if project_id else f"{langfuse_host}/sessions/{session_id}"
        except Exception as url_error:
            # Fallback if project fetch fails
            langfuse_host = config.get('langfuse_host', 'https://cloud.langfuse.com')
            session_url = f"{langfuse_host}/sessions/{session_id}"

        # Call list_traces with session_id filter
        traces_json = list_traces(session_id=session_id, include_observations=False)

        if json_output:
            # Return raw JSON
            traces_data = json.loads(traces_json) if isinstance(traces_json, str) else traces_json
            return {
                "success": True,
                "session_id": session_id,
                "session_url": session_url,
                "traces": traces_data
            }
        else:
            # Return formatted table
            table_output = format_traces_table(traces_json)
            return {
                "success": True,
                "session_id": session_id,
                "session_url": session_url,
                "table": table_output
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Error viewing session traces: {str(e)}"
        }


# ============================================================================
# Langfuse Prompts Tools
# ============================================================================

async def coaia_fuse_prompts_list() -> Dict[str, Any]:
    """
    List all Langfuse prompts.
    
    Returns:
        Dict with success status and list of prompts/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's list_prompts function
        prompts_data = cofuse_list_prompts(debug=False)
        
        # Parse the response (it might be a formatted string or dict)
        if isinstance(prompts_data, str):
            # If it's a formatted string, we return it as-is
            return {
                "success": True,
                "prompts": prompts_data,
                "note": "Prompts returned as formatted string"
            }
        elif isinstance(prompts_data, (list, dict)):
            return {
                "success": True,
                "prompts": prompts_data
            }
        else:
            return {
                "success": True,
                "prompts": str(prompts_data)
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse prompts list error: {str(e)}"
        }


async def coaia_fuse_prompts_get(name: str, label: Optional[str] = None, version: Optional[int] = None) -> Dict[str, Any]:
    """
    Get specific Langfuse prompt by name, with optional version or label filtering.
    
    Args:
        name: Prompt name
        label: Optional prompt label (e.g., "production", "latest")
        version: Optional integer version number for a specific prompt version
        
    Returns:
        Dict with success status and prompt data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's get_prompt function
        prompt_data = cofuse_get_prompt(prompt_name=name, label=label, version=version)
        
        return {
            "success": True,
            "prompt": prompt_data
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse prompt get error: {str(e)}"
        }


async def coaia_fuse_prompts_create(
    name: str,
    content: str,
    prompt_type: str = "text",
    labels: Optional[List[str]] = None,
    tags: Optional[List[str]] = None,
    commit_message: Optional[str] = None,
    config: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Create a new prompt version in Langfuse.
    
    Args:
        name: Prompt name
        content: Prompt content (string for text, JSON string for chat)
        prompt_type: "text" or "chat"
        labels: Optional deployment labels (e.g. ["production", "staging"])
        tags: Optional tags for categorization
        commit_message: Optional version commit message
        config: Optional config object (e.g. temperature, model settings)
        
    Returns:
        Dict with success status and created prompt data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        result = cofuse_create_prompt(
            prompt_name=name,
            content=content,
            prompt_type=prompt_type,
            labels=labels,
            tags=tags,
            commit_message=commit_message,
            config=config,
        )
        return {
            "success": True,
            "prompt": result
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse prompt create error: {str(e)}"
        }


async def coaia_fuse_prompt_version_labels_update(
    name: str,
    version: int,
    labels: List[str],
) -> Dict[str, Any]:
    """
    Update labels for a specific prompt version in Langfuse.
    Labels are unique across versions. The "latest" label is reserved.
    
    Args:
        name: Prompt name
        version: Version number to update
        labels: New labels to assign to this version
        
    Returns:
        Dict with success status and updated prompt data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        result = cofuse_update_prompt_version_labels(
            prompt_name=name,
            version=version,
            new_labels=labels,
        )
        return {
            "success": True,
            "prompt": result
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse prompt version labels update error: {str(e)}"
        }


# ============================================================================
# Langfuse Datasets Tools
# ============================================================================

async def coaia_fuse_datasets_list() -> Dict[str, Any]:
    """
    List all Langfuse datasets.
    
    Returns:
        Dict with success status and list of datasets/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's list_datasets function
        datasets_data = cofuse_list_datasets()
        
        return {
            "success": True,
            "datasets": datasets_data
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse datasets list error: {str(e)}"
        }


async def coaia_fuse_datasets_get(name: str) -> Dict[str, Any]:
    """
    Get specific Langfuse dataset.
    
    Args:
        name: Dataset name
        
    Returns:
        Dict with success status and dataset data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's get_dataset function
        dataset_data = cofuse_get_dataset(dataset_name=name)
        
        return {
            "success": True,
            "dataset": dataset_data
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse dataset get error: {str(e)}"
        }


# ============================================================================
# Langfuse Score Configurations Tools
# ============================================================================

async def coaia_fuse_score_configs_list() -> Dict[str, Any]:
    """
    List all Langfuse score configurations using coaiapy's smart cache system.
    
    Returns:
        Dict with success status and list of configs/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's list_score_configs function
        configs_data = list_score_configs(debug=False)
        
        return {
            "success": True,
            "configs": configs_data
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse score configs list error: {str(e)}"
        }


async def coaia_fuse_score_configs_get(name_or_id: str) -> Dict[str, Any]:
    """
    Get specific score configuration using coaiapy's smart cache system.
    
    Args:
        name_or_id: Score config name or ID
        
    Returns:
        Dict with success status and config data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's get_score_config function
        config_data = get_score_config(config_id=name_or_id)
        
        return {
            "success": True,
            "config": config_data
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse score config get error: {str(e)}"
        }


async def coaia_fuse_score_apply(
    config_name_or_id: str,
    target_type: str,
    target_id: str,
    value: Any,
    observation_id: Optional[str] = None,
    comment: Optional[str] = None
) -> Dict[str, Any]:
    """
    Apply a score configuration to a trace or observation with validation.
    
    Args:
        config_name_or_id: Name or ID of the score configuration to apply
        target_type: Type of target - "trace" or "session"
        target_id: ID of the trace or session
        value: Score value to apply (will be validated against config constraints)
        observation_id: Optional observation ID (only for trace targets)
        comment: Optional comment to attach to the score
        
    Returns:
        Dict with success status and score application result/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }
    
    try:
        # Use coaiapy's apply_score_config function which handles validation
        result = apply_score_config(
            config_name_or_id=config_name_or_id,
            target_type=target_type,
            target_id=target_id,
            value=value,
            observation_id=observation_id,
            comment=comment
        )
        
        # Check if result indicates an error
        if isinstance(result, str) and result.startswith("Error:"):
            return {
                "success": False,
                "error": result
            }
        
        return {
            "success": True,
            "message": f"Score config '{config_name_or_id}' applied to {target_type} '{target_id}'",
            "target_type": target_type,
            "target_id": target_id,
            "observation_id": observation_id,
            "value": value,
            "comment": comment,
            "result": result
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Score application error: {str(e)}"
        }


# ============================================================================
# Langfuse Comments Tools
# ============================================================================

async def coaia_fuse_comments_list(
    object_type: Optional[str] = None,
    object_id: Optional[str] = None,
    author_user_id: Optional[str] = None,
    page: int = 1,
    limit: int = 50
) -> Dict[str, Any]:
    """
    List comments with optional filtering.

    Args:
        object_type: Filter by object type (trace, observation, session, prompt)
        object_id: Filter by specific object ID (requires object_type)
        author_user_id: Filter by author user ID
        page: Page number (starts at 1)
        limit: Items per page

    Returns:
        Dict with success status and list of comments/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }

    try:
        # Use coaiapy's get_comments function
        comments_data = get_comments(
            object_type=object_type,
            object_id=object_id,
            author_user_id=author_user_id,
            page=page,
            limit=limit
        )

        # Parse response if it's a string
        import json
        if isinstance(comments_data, str):
            try:
                parsed_data = json.loads(comments_data)
                return {
                    "success": True,
                    "comments": parsed_data
                }
            except json.JSONDecodeError:
                return {
                    "success": True,
                    "comments": comments_data,
                    "note": "Comments returned as raw string"
                }
        else:
            return {
                "success": True,
                "comments": comments_data
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse comments list error: {str(e)}"
        }


async def coaia_fuse_comments_get(comment_id: str) -> Dict[str, Any]:
    """
    Get a specific comment by ID.

    Args:
        comment_id: The unique Langfuse identifier of a comment

    Returns:
        Dict with success status and comment data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }

    try:
        # Use coaiapy's get_comment_by_id function
        comment_data = get_comment_by_id(comment_id)

        # Parse response if it's a string
        import json
        if isinstance(comment_data, str):
            try:
                parsed_data = json.loads(comment_data)
                return {
                    "success": True,
                    "comment": parsed_data
                }
            except json.JSONDecodeError:
                return {
                    "success": True,
                    "comment": comment_data,
                    "note": "Comment returned as raw string"
                }
        else:
            return {
                "success": True,
                "comment": comment_data
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse comment get error: {str(e)}"
        }


async def coaia_fuse_comments_create(
    text: str,
    object_type: str,
    object_id: str,
    author_user_id: Optional[str] = None
) -> Dict[str, Any]:
    """
    Create a comment attached to an object (trace, observation, session, or prompt).

    Args:
        text: The comment text/content
        object_type: Type of object to attach comment to (trace, observation, session, prompt) - REQUIRED
        object_id: ID of the object to attach comment to - REQUIRED
        author_user_id: Optional user ID of the comment author

    Returns:
        Dict with success status and created comment data/error
    """
    if not LANGFUSE_AVAILABLE:
        return {
            "success": False,
            "error": "Langfuse is not available. Check credentials in configuration."
        }

    try:
        # Use coaiapy's post_comment function
        comment_data = post_comment(
            text=text,
            object_type=object_type,
            object_id=object_id,
            author_user_id=author_user_id
        )

        # Parse response if it's a string
        import json
        if isinstance(comment_data, str):
            try:
                parsed_data = json.loads(comment_data)
                return {
                    "success": True,
                    "comment": parsed_data
                }
            except json.JSONDecodeError:
                return {
                    "success": True,
                    "comment": comment_data,
                    "note": "Comment returned as raw string"
                }
        else:
            return {
                "success": True,
                "comment": comment_data
            }
    except Exception as e:
        return {
            "success": False,
            "error": f"Langfuse comment creation error: {str(e)}"
        }


# ============================================================================
# Langfuse Media Upload Tools
# ============================================================================

async def coaia_fuse_media_upload(
    file_path: str,
    trace_id: str,
    field: str = "input",
    observation_id: Optional[str] = None,
    content_type: Optional[str] = None,
    json_output: bool = False
) -> Dict[str, Any]:
    """
    Upload a file and attach it to a Langfuse trace or observation.

    Uploads images, videos, audio, documents (52 content types) with automatic
    MIME type detection, SHA-256 deduplication, and S3 storage. Returns media_id
    for later retrieval.

    Args:
        file_path (str): Path to file (e.g., "photo.jpg", "./docs/report.pdf")
        trace_id (str): Langfuse trace ID (e.g., "trace_abc123")
        field (str): Semantic context - "input", "output", or "metadata" (default: "input")
        observation_id (str, optional): Attach to observation instead of trace
        content_type (str, optional): MIME type override (usually auto-detected)
        json_output (bool): Return raw JSON (default: False returns formatted display)

    Returns:
        dict: {
            "success": bool,
            "media_id": str - Use this with coaia_fuse_media_get,
            "message": str - Success message with file size,
            "upload_time_ms": float - Upload duration,
            "formatted_display": str - Human-readable output (if json_output=False),
            "error": str - Error message (only if success=False)
        }

    Examples:
        Upload image to trace:
        >>> result = await coaia_fuse_media_upload(
        ...     file_path="sketch.jpg",
        ...     trace_id="trace_001",
        ...     field="input"
        ... )
        >>> print(result["media_id"])  # "media_xyz789"

        Upload audio to observation:
        >>> result = await coaia_fuse_media_upload(
        ...     file_path="recording.mp3",
        ...     trace_id="trace_001",
        ...     observation_id="obs_456",
        ...     field="output"
        ... )
    """
    try:
        result = upload_and_attach_media(
            file_path=file_path,
            trace_id=trace_id,
            field=field,
            observation_id=observation_id,
            content_type=content_type
        )

        if json_output:
            return result

        # Format friendly output
        if result["success"]:
            formatted = format_media_display(result['media_data'])
            return {
                "success": True,
                "media_id": result['media_id'],
                "message": result['message'],
                "upload_time_ms": result['upload_time_ms'],
                "formatted_display": formatted
            }
        else:
            # Propagate error details if available from underlying function
            if not result["success"] and "detail" in result:
                return {
                    "success": False,
                    "error": result["error"],
                    "detail": result["detail"]
                }
            return result

    except Exception as e:
        return {
            "success": False,
            "error": f"Media upload error: {str(e)}"
        }


async def coaia_fuse_media_get(
    media_id: str,
    json_output: bool = False
) -> Dict[str, Any]:
    """
    Retrieve metadata about a previously uploaded media file.

    Returns information about a media file including content type, size,
    trace/observation linkage, upload timestamp, and SHA-256 hash. Use the
    media_id from coaia_fuse_media_upload response.

    Args:
        media_id (str): Media ID from upload (e.g., "media_xyz789")
        json_output (bool): Return raw JSON (default: False returns formatted display)

    Returns:
        dict: {
            "success": bool,
            "media": dict - Media object with all metadata,
            "formatted_display": str - Human-readable output with icons (if json_output=False),
            "error": str - Error message (only if success=False)
        }

        Media object contains:
        - id: Media ID
        - traceId: Associated trace
        - observationId: Associated observation (if any)
        - field: "input", "output", or "metadata"
        - contentType: MIME type (e.g., "image/jpeg")
        - contentLength: File size in bytes
        - sha256Hash: Deduplication hash
        - uploadedAt: ISO timestamp

    Examples:
        Retrieve media metadata:
        >>> result = await coaia_fuse_media_get("media_xyz789")
        >>> print(result["formatted_display"])
        🖼️ Media: photo.jpg
        ├── 🆔 ID: media_xyz789
        ├── 📝 Content Type: image/jpeg
        └── 📏 Size: 193424 bytes

        Get raw data:
        >>> result = await coaia_fuse_media_get("media_xyz789", json_output=True)
        >>> print(f"Size: {result['media']['contentLength']} bytes")
    """
    try:
        import json

        media_json_str = get_media(media_id)
        media_data = json.loads(media_json_str)

        if "error" in media_data:
            return {
                "success": False,
                "error": media_data["error"]
            }

        if json_output:
            return {
                "success": True,
                "media": media_data
            }

        # Format friendly output
        formatted = format_media_display(media_data)
        return {
            "success": True,
            "media": media_data,
            "formatted_display": formatted
        }

    except Exception as e:
        return {
            "success": False,
            "error": f"Media retrieval error: {str(e)}"
        }


# ============================================================================
# Tool Registry
# ============================================================================

# Export all tools
TOOLS = {
    # Redis tools
    "coaia_tash": coaia_tash,
    "coaia_fetch": coaia_fetch,

    # Langfuse trace tools
    "coaia_fuse_trace_create": coaia_fuse_trace_create,
    "coaia_fuse_add_observation": coaia_fuse_add_observation,
    "coaia_fuse_trace_patch_output": coaia_fuse_trace_patch_output,
    "coaia_fuse_trace_get": coaia_fuse_trace_get,
    "coaia_fuse_trace_view": coaia_fuse_trace_view,
    "coaia_fuse_observation_get": coaia_fuse_observation_get,
    "coaia_fuse_traces_list": coaia_fuse_traces_list,
    "coaia_fuse_traces_session_view": coaia_fuse_traces_session_view,

    # Langfuse prompts tools
    "coaia_fuse_prompts_list": coaia_fuse_prompts_list,
    "coaia_fuse_prompts_get": coaia_fuse_prompts_get,
    "coaia_fuse_prompts_create": coaia_fuse_prompts_create,
    "coaia_fuse_prompt_version_labels_update": coaia_fuse_prompt_version_labels_update,

    # Langfuse datasets tools
    "coaia_fuse_datasets_list": coaia_fuse_datasets_list,
    "coaia_fuse_datasets_get": coaia_fuse_datasets_get,

    # Langfuse score configs tools
    "coaia_fuse_score_configs_list": coaia_fuse_score_configs_list,
    "coaia_fuse_score_configs_get": coaia_fuse_score_configs_get,
    "coaia_fuse_score_apply": coaia_fuse_score_apply,

    # Langfuse comments tools
    "coaia_fuse_comments_list": coaia_fuse_comments_list,
    "coaia_fuse_comments_get": coaia_fuse_comments_get,
    "coaia_fuse_comments_create": coaia_fuse_comments_create,

    # Langfuse media tools
    "coaia_fuse_media_upload": coaia_fuse_media_upload,
    "coaia_fuse_media_get": coaia_fuse_media_get,
}

__all__ = [
    "TOOLS",
    "coaia_tash",
    "coaia_fetch",
    "coaia_fuse_trace_create",
    "coaia_fuse_add_observation",
    "coaia_fuse_trace_patch_output",
    "coaia_fuse_trace_get",
    "coaia_fuse_trace_view",
    "coaia_fuse_observation_get",
    "coaia_fuse_traces_list",
    "coaia_fuse_traces_session_view",
    "coaia_fuse_prompts_list",
    "coaia_fuse_prompts_get",
    "coaia_fuse_prompts_create",
    "coaia_fuse_prompt_version_labels_update",
    "coaia_fuse_datasets_list",
    "coaia_fuse_datasets_get",
    "coaia_fuse_score_configs_list",
    "coaia_fuse_score_configs_get",
    "coaia_fuse_score_apply",
    "coaia_fuse_comments_list",
    "coaia_fuse_comments_get",
    "coaia_fuse_comments_create",
    "coaia_fuse_media_upload",
    "coaia_fuse_media_get",
]
