main.py

"""
DFCX SCRAPI MCP Server
======================
A comprehensive Model Control Protocol (MCP) server for Google Dialogflow CX SCRAPI
Provides high-level bot building and management capabilities through MCP tools
"""

import os
import json
import uuid
import asyncio
import base64
from typing import Any, Dict, List, Optional, Union
from pathlib import Path
from datetime import datetime
import pandas as pd
import io

# Google Cloud imports
from google.oauth2 import service_account
from google.cloud import dialogflowcx_v3 as dialogflow

# DFCX SCRAPI imports
from dfcx_scrapi.core.agents import Agents
from dfcx_scrapi.core.intents import Intents
from dfcx_scrapi.core.flows import Flows
from dfcx_scrapi.core.pages import Pages
from dfcx_scrapi.core.entity_types import EntityTypes
from dfcx_scrapi.core.webhooks import Webhooks
from dfcx_scrapi.core.transition_route_groups import TransitionRouteGroups
from dfcx_scrapi.core.sessions import Sessions
from dfcx_scrapi.core.environments import Environments
from dfcx_scrapi.core.test_cases import TestCases
from dfcx_scrapi.core.playbooks import Playbooks
from dfcx_scrapi.core.conversation import DialogflowConversation
from dfcx_scrapi.core.data_stores import DataStores
from dfcx_scrapi.core.engines import Engines
from dfcx_scrapi.core.sites import Sites
from dfcx_scrapi.core.generators import Generators

# SCRAPI Tools imports
from dfcx_scrapi.tools.dataframe_functions import DataframeFunctions
from dfcx_scrapi.tools.copy_util import CopyUtil
from dfcx_scrapi.tools.search_util import SearchUtil
from dfcx_scrapi.tools.agent_extract import AgentExtract
from dfcx_scrapi.tools.nlu_util import NLUUtil
from dfcx_scrapi.tools.stats_util import StatsUtil

# MCP Server
from mcp.server.fastmcp import FastMCP

# Initialize FastMCP server
mcp = FastMCP("dfcx-scrapi")

# Global state management
class DFCXScrapiState:
    """Manages global state for DFCX SCRAPI operations"""
    
    def __init__(self):
        self.creds = None
        self.creds_path = None
        self.project_id = None
        self.location = None
        self.agent_id = None
        
        # Core objects
        self.agents = None
        self.intents = None
        self.flows = None
        self.pages = None
        self.entity_types = None
        self.webhooks = None
        self.route_groups = None
        self.sessions = None
        self.environments = None
        self.test_cases = None
        self.playbooks = None
        self.conversation = None
        self.data_stores = None
        self.engines = None
        self.sites = None
        self.generators = None
        
        # Tools objects
        self.df_functions = None
        self.copy_util = None
        self.search_util = None
        self.agent_extract = None
        self.nlu_util = None
        self.stats_util = None
        
    def initialize_credentials(self, creds_path: str = None, project_id: str = None):
        """Initialize credentials and project settings"""
        if creds_path and os.path.exists(creds_path):
            self.creds = service_account.Credentials.from_service_account_file(creds_path)
            self.creds_path = creds_path
        else:
            # Use default credentials
            self.creds = None
            
        if project_id:
            self.project_id = project_id
            
    def initialize_core_objects(self, agent_id: str = None):
        """Initialize core SCRAPI objects"""
        if agent_id:
            self.agent_id = agent_id
            
        # Initialize core objects
        self.agents = Agents(creds=self.creds, agent_id=self.agent_id)
        self.intents = Intents(creds=self.creds, agent_id=self.agent_id)
        self.flows = Flows(creds=self.creds, agent_id=self.agent_id)
        self.pages = Pages(creds=self.creds, agent_id=self.agent_id)
        self.entity_types = EntityTypes(creds=self.creds, agent_id=self.agent_id)
        self.webhooks = Webhooks(creds=self.creds, agent_id=self.agent_id)
        self.route_groups = TransitionRouteGroups(creds=self.creds, agent_id=self.agent_id)
        self.sessions = Sessions(creds=self.creds, agent_id=self.agent_id)
        self.environments = Environments(creds=self.creds, agent_id=self.agent_id)
        self.test_cases = TestCases(creds=self.creds, agent_id=self.agent_id)
        self.playbooks = Playbooks(creds=self.creds, agent_id=self.agent_id)
        
        # Initialize tools
        self.df_functions = DataframeFunctions(creds=self.creds)
        self.copy_util = CopyUtil(creds=self.creds)
        self.search_util = SearchUtil(creds=self.creds, agent_id=self.agent_id)
        self.stats_util = StatsUtil(creds=self.creds)
        
        # Initialize Vertex AI objects if project_id is set
        if self.project_id:
            self.data_stores = DataStores(project_id=self.project_id, creds=self.creds)
            self.engines = Engines(project_id=self.project_id, creds=self.creds)
            self.sites = Sites(project_id=self.project_id, creds=self.creds)
            self.generators = Generators(creds=self.creds, agent_id=self.agent_id)

# Create global state instance
state = DFCXScrapiState()

# ============================================================================
# INITIALIZATION AND SETUP TOOLS
# ============================================================================

@mcp.tool()
async def initialize_dfcx_scrapi(
    creds_path: Optional[str] = None,
    project_id: Optional[str] = None,
    location: str = "us-central1",
    agent_id: Optional[str] = None
) -> str:
    """
    Initialize DFCX SCRAPI with credentials and project settings.
    
    Args:
        creds_path: Path to service account credentials JSON file
        project_id: Google Cloud project ID
        location: Agent location (default: us-central1)
        agent_id: Optional agent ID to set as default
        
    Returns:
        Success message with initialization status
    """
    try:
        state.location = location
        state.initialize_credentials(creds_path, project_id)
        
        if agent_id:
            state.initialize_core_objects(agent_id)
            return f"DFCX SCRAPI initialized successfully with agent: {agent_id}"
        else:
            return "DFCX SCRAPI initialized successfully. Set an agent_id to access agent-specific features."
            
    except Exception as e:
        return f"Error initializing DFCX SCRAPI: {str(e)}"

# ============================================================================
# AGENT MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_agents(location: Optional[str] = None) -> List[Dict[str, Any]]:
    """
    List all agents in the project.
    
    Args:
        location: Optional location filter
        
    Returns:
        List of agent information dictionaries
    """
    try:
        if not state.agents:
            state.agents = Agents(creds=state.creds)
            
        agents = state.agents.list_agents(location=location or state.location)
        
        return [
            {
                "id": agent.name,
                "display_name": agent.display_name,
                "default_language": agent.default_language_code,
                "time_zone": agent.time_zone,
                "description": agent.description,
                "supported_languages": list(agent.supported_language_codes)
            }
            for agent in agents
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def create_agent(
    display_name: str,
    default_language_code: str = "en",
    time_zone: str = "America/Los_Angeles",
    location: Optional[str] = None,
    description: Optional[str] = None
) -> Dict[str, Any]:
    """
    Create a new Dialogflow CX agent.
    
    Args:
        display_name: Agent display name
        default_language_code: Default language (default: en)
        time_zone: Time zone (default: America/Los_Angeles)
        location: Agent location
        description: Optional agent description
        
    Returns:
        Created agent information
    """
    try:
        if not state.agents:
            state.agents = Agents(creds=state.creds)
            
        agent = state.agents.create_agent(
            project_id=state.project_id,
            display_name=display_name,
            default_language_code=default_language_code,
            time_zone=time_zone,
            location=location or state.location,
            description=description
        )
        
        return {
            "id": agent.name,
            "display_name": agent.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def export_agent(
    agent_id: Optional[str] = None,
    gcs_bucket_uri: Optional[str] = None,
    data_format: str = "JSON",
    environment: Optional[str] = None
) -> Dict[str, Any]:
    """
    Export an agent to GCS or get export data.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        gcs_bucket_uri: Optional GCS bucket URI for export
        data_format: Export format (JSON or BLOB)
        environment: Optional environment to export
        
    Returns:
        Export operation details
    """
    try:
        if not state.agents:
            state.agents = Agents(creds=state.creds)
            
        agent_id = agent_id or state.agent_id
        
        if gcs_bucket_uri:
            operation_id = state.agents.export_agent(
                agent_id=agent_id,
                gcs_bucket_uri=gcs_bucket_uri,
                data_format=data_format,
                environment=environment
            )
            return {"operation_id": operation_id, "status": "exporting"}
        else:
            # Export to bytes
            agent_bytes = state.agents.export_agent(
                agent_id=agent_id,
                data_format=data_format,
                environment=environment
            )
            return {
                "data": base64.b64encode(agent_bytes).decode('utf-8'),
                "format": data_format
            }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def restore_agent(
    agent_id: Optional[str] = None,
    gcs_bucket_uri: Optional[str] = None,
    agent_bytes: Optional[str] = None
) -> Dict[str, Any]:
    """
    Restore an agent from GCS or bytes data.
    
    Args:
        agent_id: Agent ID to restore to
        gcs_bucket_uri: GCS bucket URI with agent export
        agent_bytes: Base64 encoded agent data
        
    Returns:
        Restore operation details
    """
    try:
        if not state.agents:
            state.agents = Agents(creds=state.creds)
            
        agent_id = agent_id or state.agent_id
        
        if gcs_bucket_uri:
            operation_id = state.agents.restore_agent(
                agent_id=agent_id,
                gcs_bucket_uri=gcs_bucket_uri
            )
        elif agent_bytes:
            bytes_data = base64.b64decode(agent_bytes)
            operation_id = state.agents.restore_agent(
                agent_id=agent_id,
                agent_bytes=bytes_data
            )
        else:
            return {"error": "Either gcs_bucket_uri or agent_bytes must be provided"}
            
        return {"operation_id": operation_id, "status": "restoring"}
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# INTENT MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_intents(
    agent_id: Optional[str] = None,
    language_code: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all intents in an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        language_code: Optional language code filter
        
    Returns:
        List of intent information
    """
    try:
        if not state.intents:
            state.intents = Intents(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        intents = state.intents.list_intents(agent_id, language_code)
        
        return [
            {
                "id": intent.name,
                "display_name": intent.display_name,
                "priority": intent.priority,
                "is_fallback": intent.is_fallback,
                "training_phrases_count": len(intent.training_phrases),
                "parameters_count": len(intent.parameters)
            }
            for intent in intents
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def create_intent(
    display_name: str,
    training_phrases: List[str],
    agent_id: Optional[str] = None,
    parameters: Optional[List[Dict[str, Any]]] = None,
    responses: Optional[List[str]] = None
) -> Dict[str, Any]:
    """
    Create a new intent.
    
    Args:
        display_name: Intent display name
        training_phrases: List of training phrase texts
        agent_id: Agent ID (uses default if not provided)
        parameters: Optional list of parameter definitions
        responses: Optional list of response messages
        
    Returns:
        Created intent information
    """
    try:
        if not state.intents:
            state.intents = Intents(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        # Build training phrases
        tp_list = []
        for phrase in training_phrases:
            tp = state.intents.build_training_phrase(phrase)
            tp_list.append(tp)
            
        # Create intent
        intent = state.intents.create_intent(
            agent_id=agent_id,
            display_name=display_name,
            training_phrases=tp_list
        )
        
        return {
            "id": intent.name,
            "display_name": intent.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def bulk_update_intents_from_dataframe(
    csv_data: str,
    agent_id: Optional[str] = None,
    update_flag: bool = False
) -> Dict[str, Any]:
    """
    Bulk update intents from CSV data.
    
    Args:
        csv_data: CSV string with intent data
        agent_id: Agent ID (uses default if not provided)
        update_flag: Whether to update existing intents
        
    Returns:
        Update operation results
    """
    try:
        if not state.intents:
            state.intents = Intents(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        # Parse CSV to DataFrame
        df = pd.read_csv(io.StringIO(csv_data))
        
        # Bulk update
        results = state.intents.bulk_update_intents_from_dataframe(
            agent_id=agent_id,
            df=df,
            update_flag=update_flag
        )
        
        return {
            "updated_count": len(results),
            "status": "completed"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# FLOW MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_flows(
    agent_id: Optional[str] = None,
    language_code: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all flows in an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        language_code: Optional language code filter
        
    Returns:
        List of flow information
    """
    try:
        if not state.flows:
            state.flows = Flows(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        flows = state.flows.list_flows(agent_id, language_code)
        
        return [
            {
                "id": flow.name,
                "display_name": flow.display_name,
                "description": flow.description,
                "transition_routes_count": len(flow.transition_routes),
                "event_handlers_count": len(flow.event_handlers)
            }
            for flow in flows
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def create_flow(
    display_name: str,
    agent_id: Optional[str] = None,
    description: Optional[str] = None
) -> Dict[str, Any]:
    """
    Create a new flow.
    
    Args:
        display_name: Flow display name
        agent_id: Agent ID (uses default if not provided)
        description: Optional flow description
        
    Returns:
        Created flow information
    """
    try:
        if not state.flows:
            state.flows = Flows(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        flow = state.flows.create_flow(
            agent_id=agent_id,
            display_name=display_name,
            description=description
        )
        
        return {
            "id": flow.name,
            "display_name": flow.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# PAGE MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_pages(
    flow_id: str,
    language_code: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all pages in a flow.
    
    Args:
        flow_id: Flow ID
        language_code: Optional language code filter
        
    Returns:
        List of page information
    """
    try:
        if not state.pages:
            state.pages = Pages(creds=state.creds)
            
        pages = state.pages.list_pages(flow_id, language_code)
        
        return [
            {
                "id": page.name,
                "display_name": page.display_name,
                "entry_fulfillment": bool(page.entry_fulfillment),
                "form_parameters_count": len(page.form.parameters) if page.form else 0,
                "transition_routes_count": len(page.transition_routes),
                "event_handlers_count": len(page.event_handlers)
            }
            for page in pages
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def create_page(
    flow_id: str,
    display_name: str,
    entry_fulfillment_text: Optional[str] = None
) -> Dict[str, Any]:
    """
    Create a new page in a flow.
    
    Args:
        flow_id: Flow ID
        display_name: Page display name
        entry_fulfillment_text: Optional entry fulfillment message
        
    Returns:
        Created page information
    """
    try:
        if not state.pages:
            state.pages = Pages(creds=state.creds)
            
        page = state.pages.create_page(
            flow_id=flow_id,
            display_name=display_name
        )
        
        # Add entry fulfillment if provided
        if entry_fulfillment_text:
            page.entry_fulfillment = state.pages.build_fulfillment(
                messages=[{"text": entry_fulfillment_text}]
            )
            page = state.pages.update_page(page.name, page)
        
        return {
            "id": page.name,
            "display_name": page.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# ENTITY TYPE MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_entity_types(
    agent_id: Optional[str] = None,
    language_code: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all entity types in an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        language_code: Optional language code filter
        
    Returns:
        List of entity type information
    """
    try:
        if not state.entity_types:
            state.entity_types = EntityTypes(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        entity_types = state.entity_types.list_entity_types(agent_id, language_code)
        
        return [
            {
                "id": et.name,
                "display_name": et.display_name,
                "kind": et.kind.name,
                "entities_count": len(et.entities),
                "auto_expansion": et.auto_expansion_mode.name
            }
            for et in entity_types
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def create_entity_type(
    display_name: str,
    entities: List[Dict[str, List[str]]],
    agent_id: Optional[str] = None,
    kind: str = "KIND_MAP"
) -> Dict[str, Any]:
    """
    Create a new entity type.
    
    Args:
        display_name: Entity type display name
        entities: List of entities with values and synonyms
        agent_id: Agent ID (uses default if not provided)
        kind: Entity type kind (KIND_MAP or KIND_LIST)
        
    Returns:
        Created entity type information
    """
    try:
        if not state.entity_types:
            state.entity_types = EntityTypes(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        # Build entities
        entity_list = []
        for entity_data in entities:
            entity = state.entity_types.build_entity(
                value=entity_data.get("value"),
                synonyms=entity_data.get("synonyms", [])
            )
            entity_list.append(entity)
            
        # Create entity type
        entity_type = state.entity_types.create_entity_type(
            agent_id=agent_id,
            display_name=display_name,
            entities=entity_list,
            kind=kind
        )
        
        return {
            "id": entity_type.name,
            "display_name": entity_type.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# CONVERSATION AND TESTING TOOLS
# ============================================================================

@mcp.tool()
async def detect_intent_text(
    text: str,
    session_id: Optional[str] = None,
    agent_id: Optional[str] = None,
    language_code: str = "en"
) -> Dict[str, Any]:
    """
    Detect intent from text input.
    
    Args:
        text: User input text
        session_id: Session ID (generates new if not provided)
        agent_id: Agent ID (uses default if not provided)
        language_code: Language code (default: en)
        
    Returns:
        Intent detection results
    """
    try:
        agent_id = agent_id or state.agent_id
        session_id = session_id or str(uuid.uuid4())
        
        if not state.sessions:
            state.sessions = Sessions(creds=state.creds, agent_id=agent_id)
            
        response = state.sessions.detect_intent(
            agent_id=agent_id,
            session_id=session_id,
            text=text,
            language_code=language_code
        )
        
        return {
            "session_id": session_id,
            "query_text": response.query_result.text,
            "response_messages": [
                msg.text.text[0] if msg.text else str(msg)
                for msg in response.query_result.response_messages
            ],
            "intent": response.query_result.match.intent.display_name if response.query_result.match.intent else None,
            "confidence": response.query_result.match.confidence,
            "parameters": dict(response.query_result.parameters) if response.query_result.parameters else {},
            "current_page": response.query_result.current_page.display_name if response.query_result.current_page else None
        }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def run_conversation(
    test_cases: List[Dict[str, str]],
    agent_id: Optional[str] = None,
    language_code: str = "en"
) -> List[Dict[str, Any]]:
    """
    Run a full conversation with multiple turns.
    
    Args:
        test_cases: List of test cases with user inputs
        agent_id: Agent ID (uses default if not provided)
        language_code: Language code (default: en)
        
    Returns:
        Conversation results for each turn
    """
    try:
        agent_id = agent_id or state.agent_id
        
        if not state.conversation:
            state.conversation = DialogflowConversation(
                creds=state.creds,
                agent_id=agent_id
            )
            
        # Convert test cases to DataFrame
        df = pd.DataFrame(test_cases)
        
        # Run conversation
        results = state.conversation.run_conversation(
            df=df,
            agent_id=agent_id,
            language_code=language_code
        )
        
        return results.to_dict('records')
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def run_intent_detection(
    test_set: List[Dict[str, str]],
    agent_id: Optional[str] = None,
    language_code: str = "en"
) -> List[Dict[str, Any]]:
    """
    Run batch intent detection tests.
    
    Args:
        test_set: List of test cases with utterances
        agent_id: Agent ID (uses default if not provided)
        language_code: Language code (default: en)
        
    Returns:
        Intent detection results for each test case
    """
    try:
        agent_id = agent_id or state.agent_id
        
        if not state.conversation:
            state.conversation = DialogflowConversation(
                creds=state.creds,
                agent_id=agent_id
            )
            
        # Convert to DataFrame
        df = pd.DataFrame(test_set)
        
        # Run intent detection
        results = state.conversation.run_intent_detection(
            df=df,
            agent_id=agent_id,
            language_code=language_code
        )
        
        return results.to_dict('records')
    except Exception as e:
        return [{"error": str(e)}]

# ============================================================================
# WEBHOOK MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_webhooks(
    agent_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all webhooks in an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        List of webhook information
    """
    try:
        if not state.webhooks:
            state.webhooks = Webhooks(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        webhooks = state.webhooks.list_webhooks(agent_id)
        
        return [
            {
                "id": webhook.name,
                "display_name": webhook.display_name,
                "webhook_uri": webhook.generic_web_service.uri if webhook.generic_web_service else None,
                "timeout": webhook.timeout.seconds if webhook.timeout else None,
                "disabled": webhook.disabled
            }
            for webhook in webhooks
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def create_webhook(
    display_name: str,
    webhook_uri: str,
    agent_id: Optional[str] = None,
    timeout: int = 30,
    headers: Optional[Dict[str, str]] = None
) -> Dict[str, Any]:
    """
    Create a new webhook.
    
    Args:
        display_name: Webhook display name
        webhook_uri: Webhook URI
        agent_id: Agent ID (uses default if not provided)
        timeout: Timeout in seconds (default: 30)
        headers: Optional webhook headers
        
    Returns:
        Created webhook information
    """
    try:
        if not state.webhooks:
            state.webhooks = Webhooks(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        webhook = state.webhooks.create_webhook(
            agent_id=agent_id,
            display_name=display_name,
            webhook_uri=webhook_uri,
            timeout=timeout,
            headers=headers
        )
        
        return {
            "id": webhook.name,
            "display_name": webhook.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# ENVIRONMENT MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_environments(
    agent_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all environments in an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        List of environment information
    """
    try:
        if not state.environments:
            state.environments = Environments(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        environments = state.environments.list_environments(agent_id)
        
        return [
            {
                "id": env.name,
                "display_name": env.display_name,
                "description": env.description,
                "version_configs": [
                    {"version": vc.version} for vc in env.version_configs
                ]
            }
            for env in environments
        ]
    except Exception as e:
        return [{"error": str(e)}]

# ============================================================================
# VERTEX AI SEARCH & CONVERSATION TOOLS
# ============================================================================

@mcp.tool()
async def create_data_store(
    display_name: str,
    data_store_type: str = "PUBLIC_WEB",
    solution_types: Optional[List[str]] = None
) -> Dict[str, Any]:
    """
    Create a Vertex AI Search data store.
    
    Args:
        display_name: Data store display name
        data_store_type: Type of data store (PUBLIC_WEB, UNSTRUCTURED, etc.)
        solution_types: Solution types (SOLUTION_TYPE_SEARCH, SOLUTION_TYPE_CHAT)
        
    Returns:
        Created data store information
    """
    try:
        if not state.data_stores:
            if not state.project_id:
                return {"error": "Project ID not set. Initialize with project_id first."}
            state.data_stores = DataStores(project_id=state.project_id, creds=state.creds)
            
        data_store = state.data_stores.create_datastore(
            display_name=display_name,
            data_store_type=data_store_type,
            solution_types=solution_types or ["SOLUTION_TYPE_CHAT"]
        )
        
        return {
            "id": data_store.name,
            "display_name": data_store.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def create_search_engine(
    display_name: str,
    data_store_ids: List[str],
    business_name: Optional[str] = None,
    search_tier: str = "SEARCH_TIER_STANDARD"
) -> Dict[str, Any]:
    """
    Create a Vertex AI Search engine.
    
    Args:
        display_name: Engine display name
        data_store_ids: List of data store IDs to include
        business_name: Optional business name
        search_tier: Search tier (STANDARD or ENTERPRISE)
        
    Returns:
        Created engine information
    """
    try:
        if not state.engines:
            if not state.project_id:
                return {"error": "Project ID not set. Initialize with project_id first."}
            state.engines = Engines(project_id=state.project_id, creds=state.creds)
            
        # Build engine proto
        engine_proto = state.engines.build_chat_engine_proto(
            display_name=display_name,
            data_store_ids=data_store_ids,
            business_name=business_name,
            search_tier=search_tier
        )
        
        # Create engine
        engine = state.engines.create_engine(engine_proto)
        
        return {
            "id": engine.name,
            "display_name": engine.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# GENERATOR TOOLS (VERTEX AI INTEGRATION)
# ============================================================================

@mcp.tool()
async def list_generators(
    agent_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all generators in an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        List of generator information
    """
    try:
        if not state.generators:
            state.generators = Generators(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        generators = state.generators.list_generators(agent_id)
        
        return [
            {
                "id": gen.name,
                "display_name": gen.display_name,
                "prompt_text": gen.prompt_text.text if gen.prompt_text else None,
                "llm_model": gen.llm_model_settings.model if gen.llm_model_settings else None
            }
            for gen in generators
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def create_generator(
    display_name: str,
    prompt_text: str,
    agent_id: Optional[str] = None,
    model_name: str = "gemini-1.5-flash"
) -> Dict[str, Any]:
    """
    Create a new generator with Vertex AI model.
    
    Args:
        display_name: Generator display name
        prompt_text: Prompt text for the generator
        agent_id: Agent ID (uses default if not provided)
        model_name: Model name (default: gemini-1.5-flash)
        
    Returns:
        Created generator information
    """
    try:
        if not state.generators:
            state.generators = Generators(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        generator = state.generators.create_generator(
            agent_id=agent_id,
            display_name=display_name,
            prompt_text=prompt_text,
            model_name=model_name
        )
        
        return {
            "id": generator.name,
            "display_name": generator.display_name,
            "status": "created"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# DATAFRAME AND ANALYSIS TOOLS
# ============================================================================

@mcp.tool()
async def intents_to_dataframe(
    agent_id: Optional[str] = None
) -> str:
    """
    Export all intents to a CSV string.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        CSV string of intents data
    """
    try:
        if not state.df_functions:
            state.df_functions = DataframeFunctions(creds=state.creds)
            
        agent_id = agent_id or state.agent_id
        
        # Get intents as DataFrame
        df = state.df_functions.intents_to_df(agent_id)
        
        # Convert to CSV string
        return df.to_csv(index=False)
    except Exception as e:
        return f"Error: {str(e)}"

@mcp.tool()
async def flows_to_dataframe(
    agent_id: Optional[str] = None
) -> str:
    """
    Export all flows to a CSV string.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        CSV string of flows data
    """
    try:
        if not state.df_functions:
            state.df_functions = DataframeFunctions(creds=state.creds)
            
        agent_id = agent_id or state.agent_id
        
        # Get flows as DataFrame
        df = state.df_functions.flows_to_df(agent_id)
        
        # Convert to CSV string
        return df.to_csv(index=False)
    except Exception as e:
        return f"Error: {str(e)}"

@mcp.tool()
async def pages_to_dataframe(
    flow_id: str
) -> str:
    """
    Export all pages in a flow to a CSV string.
    
    Args:
        flow_id: Flow ID
        
    Returns:
        CSV string of pages data
    """
    try:
        if not state.df_functions:
            state.df_functions = DataframeFunctions(creds=state.creds)
            
        # Get pages as DataFrame
        df = state.df_functions.pages_to_df(flow_id)
        
        # Convert to CSV string
        return df.to_csv(index=False)
    except Exception as e:
        return f"Error: {str(e)}"

# ============================================================================
# SEARCH AND UTILITY TOOLS
# ============================================================================

@mcp.tool()
async def search_agent_resources(
    query: str,
    resource_types: Optional[List[str]] = None,
    agent_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    Search across agent resources (intents, flows, pages, etc.).
    
    Args:
        query: Search query string
        resource_types: Types to search (intents, flows, pages, etc.)
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        Search results across resources
    """
    try:
        if not state.search_util:
            state.search_util = SearchUtil(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        # Search across resources
        results = state.search_util.search_all_resources(
            agent_id=agent_id,
            query=query,
            resource_types=resource_types
        )
        
        return results
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def copy_intent(
    source_intent_id: str,
    destination_agent_id: str,
    display_name: Optional[str] = None
) -> Dict[str, Any]:
    """
    Copy an intent to another agent.
    
    Args:
        source_intent_id: Source intent ID
        destination_agent_id: Destination agent ID
        display_name: Optional new display name
        
    Returns:
        Copied intent information
    """
    try:
        if not state.copy_util:
            state.copy_util = CopyUtil(creds=state.creds)
            
        result = state.copy_util.copy_intent(
            source_intent_id=source_intent_id,
            destination_agent_id=destination_agent_id,
            display_name=display_name
        )
        
        return {
            "id": result.name,
            "display_name": result.display_name,
            "status": "copied"
        }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def copy_flow(
    source_flow_id: str,
    destination_agent_id: str,
    display_name: Optional[str] = None
) -> Dict[str, Any]:
    """
    Copy a flow to another agent.
    
    Args:
        source_flow_id: Source flow ID
        destination_agent_id: Destination agent ID
        display_name: Optional new display name
        
    Returns:
        Copied flow information
    """
    try:
        if not state.copy_util:
            state.copy_util = CopyUtil(creds=state.creds)
            
        result = state.copy_util.copy_flow(
            source_flow_id=source_flow_id,
            destination_agent_id=destination_agent_id,
            display_name=display_name
        )
        
        return {
            "id": result.name,
            "display_name": result.display_name,
            "status": "copied"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# AGENT STATISTICS AND ANALYSIS TOOLS
# ============================================================================

@mcp.tool()
async def get_agent_stats(
    agent_id: Optional[str] = None
) -> Dict[str, Any]:
    """
    Get comprehensive statistics about an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        Agent statistics including counts and complexity metrics
    """
    try:
        if not state.stats_util:
            state.stats_util = StatsUtil(creds=state.creds)
            
        agent_id = agent_id or state.agent_id
        
        stats = state.stats_util.get_agent_stats(agent_id)
        
        return stats
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def extract_agent_offline(
    agent_id: Optional[str] = None
) -> Dict[str, Any]:
    """
    Extract entire agent structure offline for analysis.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        Complete agent structure with all resources
    """
    try:
        if not state.agent_extract:
            state.agent_extract = AgentExtract(creds=state.creds)
            
        agent_id = agent_id or state.agent_id
        
        # Extract agent
        agent_data = state.agent_extract.extract_agent(agent_id)
        
        return {
            "agent": agent_data.agent_dict,
            "intents": len(agent_data.intents_list),
            "flows": len(agent_data.flows_list),
            "pages": len(agent_data.pages_list),
            "entity_types": len(agent_data.entity_types_list),
            "webhooks": len(agent_data.webhooks_list),
            "routes": len(agent_data.routes_list),
            "status": "extracted"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# NLU EVALUATION TOOLS
# ============================================================================

@mcp.tool()
async def evaluate_nlu(
    test_dataset: List[Dict[str, str]],
    agent_id: Optional[str] = None,
    confidence_threshold: float = 0.3
) -> Dict[str, Any]:
    """
    Evaluate NLU performance with test dataset.
    
    Args:
        test_dataset: List of test cases with expected intents
        agent_id: Agent ID (uses default if not provided)
        confidence_threshold: Confidence threshold for predictions
        
    Returns:
        NLU evaluation metrics and results
    """
    try:
        if not state.nlu_util:
            state.nlu_util = NLUUtil(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        # Convert to DataFrame
        df = pd.DataFrame(test_dataset)
        
        # Run evaluation
        results = state.nlu_util.run_nlu_evaluation(
            agent_id=agent_id,
            test_df=df,
            confidence_threshold=confidence_threshold
        )
        
        return {
            "accuracy": results["accuracy"],
            "precision": results["precision"],
            "recall": results["recall"],
            "f1_score": results["f1_score"],
            "confusion_matrix": results["confusion_matrix"],
            "detailed_results": results["detailed_results"].to_dict('records')
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# TEST CASE MANAGEMENT TOOLS
# ============================================================================

@mcp.tool()
async def list_test_cases(
    agent_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    List all test cases in an agent.
    
    Args:
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        List of test case information
    """
    try:
        if not state.test_cases:
            state.test_cases = TestCases(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        test_cases = state.test_cases.list_test_cases(agent_id)
        
        return [
            {
                "id": tc.name,
                "display_name": tc.display_name,
                "tags": list(tc.tags),
                "test_config": {
                    "flow": tc.test_config.flow if tc.test_config else None,
                    "page": tc.test_config.page if tc.test_config else None
                }
            }
            for tc in test_cases
        ]
    except Exception as e:
        return [{"error": str(e)}]

@mcp.tool()
async def run_test_cases(
    test_case_ids: Optional[List[str]] = None,
    agent_id: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    Run test cases and get results.
    
    Args:
        test_case_ids: Optional list of test case IDs to run
        agent_id: Agent ID (uses default if not provided)
        
    Returns:
        Test case results
    """
    try:
        if not state.test_cases:
            state.test_cases = TestCases(creds=state.creds, agent_id=state.agent_id)
            
        agent_id = agent_id or state.agent_id
        
        # Run test cases
        results = state.test_cases.run_test_cases(
            agent_id=agent_id,
            test_case_ids=test_case_ids
        )
        
        return [
            {
                "test_case_id": result.name,
                "status": result.test_result.name,
                "conversation_turns": len(result.conversation_turns)
            }
            for result in results
        ]
    except Exception as e:
        return [{"error": str(e)}]

# ============================================================================
# BATCH OPERATIONS TOOLS
# ============================================================================

@mcp.tool()
async def batch_delete_intents(
    intent_ids: List[str]
) -> Dict[str, Any]:
    """
    Batch delete multiple intents.
    
    Args:
        intent_ids: List of intent IDs to delete
        
    Returns:
        Deletion results
    """
    try:
        if not state.intents:
            state.intents = Intents(creds=state.creds)
            
        deleted_count = 0
        errors = []
        
        for intent_id in intent_ids:
            try:
                state.intents.delete_intent(intent_id)
                deleted_count += 1
            except Exception as e:
                errors.append({
                    "intent_id": intent_id,
                    "error": str(e)
                })
                
        return {
            "deleted_count": deleted_count,
            "errors": errors,
            "status": "completed"
        }
    except Exception as e:
        return {"error": str(e)}

@mcp.tool()
async def batch_update_training_phrases(
    updates: List[Dict[str, Any]]
) -> Dict[str, Any]:
    """
    Batch update training phrases for multiple intents.
    
    Args:
        updates: List of updates with intent_id and training_phrases
        
    Returns:
        Update results
    """
    try:
        if not state.intents:
            state.intents = Intents(creds=state.creds)
            
        updated_count = 0
        errors = []
        
        for update in updates:
            try:
                intent_id = update["intent_id"]
                training_phrases = update["training_phrases"]
                
                # Get intent
                intent = state.intents.get_intent(intent_id)
                
                # Build new training phrases
                tp_list = []
                for phrase in training_phrases:
                    tp = state.intents.build_training_phrase(phrase)
                    tp_list.append(tp)
                    
                # Update intent
                intent.training_phrases = tp_list
                state.intents.update_intent(intent_id, intent)
                updated_count += 1
                
            except Exception as e:
                errors.append({
                    "intent_id": update.get("intent_id"),
                    "error": str(e)
                })
                
        return {
            "updated_count": updated_count,
            "errors": errors,
            "status": "completed"
        }
    except Exception as e:
        return {"error": str(e)}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

if __name__ == "__main__":
    # Print startup message
    print("DFCX SCRAPI MCP Server starting...", flush=True)
    print("This server provides comprehensive Dialogflow CX management capabilities", flush=True)
    print("Initialize with credentials and agent ID to get started", flush=True)
    
    # Run the server
    mcp.run(transport='stdio')