"""Enterprise platform meta-tools for MCP Grid.

These tools extend the base 9 meta-tools with enterprise AI platform
integration capabilities. They allow agents to:
  - Connect to enterprise platforms
  - Browse and search workflow chains
  - Configure and execute workflows
  - Get execution results with full tool call traces

All sensitive data (tokens, credentials) is held in memory only
and never logged or persisted.
"""

from __future__ import annotations

import json
import logging
from typing import Annotated, Any

from pydantic import Field

from mcp_bridge.platforms.auth.token_manager import TokenManager
from mcp_bridge.platforms.models.workflow import PlatformConfig, PlatformConnection
from mcp_bridge.platforms.workflow.executor import WorkflowExecutor

logger = logging.getLogger(__name__)

# Global references -- set during initialization
_platforms: dict[str, PlatformConnection] = {}
_executors: dict[str, WorkflowExecutor] = {}
_token_manager: TokenManager | None = None


def init_platform_tools(token_manager: TokenManager | None = None) -> None:
    """Initialize platform tools with dependencies."""
    global _token_manager
    _token_manager = token_manager or TokenManager()


def _get_token_manager() -> TokenManager:
    """Get or create the global token manager."""
    global _token_manager
    if _token_manager is None:
        _token_manager = TokenManager()
    return _token_manager


def _get_executor(platform_id: str) -> WorkflowExecutor:
    """Get the executor for a connected platform."""
    if platform_id not in _executors:
        raise RuntimeError(
            f"Platform '{platform_id}' is not connected. "
            "Use connect_platform first."
        )
    return _executors[platform_id]


# ---------------------------------------------------------------------------
# Meta-Tools (registered with FastMCP in server.py)
# ---------------------------------------------------------------------------


async def connect_platform(
    platform_id: Annotated[str, Field(description="A unique identifier for this connection (e.g., 'acme-ai').")],
    name: Annotated[str, Field(description="Display name (e.g., 'Acme Corp AI Platform').")],
    base_url: Annotated[str, Field(description="Platform frontend URL (e.g., 'https://platform.acme.com').")],
    api_base_url: Annotated[str, Field(description="REST API base (e.g., 'https://api.acme.com').")],
    websocket_url: Annotated[
        str,
        Field(
            description="WebSocket gateway base (e.g., 'wss://ws.acme.com/prod'). The token will be appended as ?Authorization=<token>."
        ),
    ],
    platform_token: Annotated[
        str,
        Field(description="JWT token from the platform's SSO authentication. This is stored in memory only."),
    ],
    heartbeat_interval: Annotated[
        int,
        Field(description="Seconds between heartbeats during execution (default 25)."),
    ] = 25,
) -> str:
    """Connect to an enterprise AI workflow platform.

    Registers a platform connection with the provided JWT token.
    The token is held in memory only and never written to disk.

    After connecting, use list_workflows to browse available chains,
    and execute_workflow to run them.

    Args:
        platform_id: A unique identifier for this connection (e.g., 'acme-ai').
        name: Display name (e.g., 'Acme Corp AI Platform').
        base_url: Platform frontend URL (e.g., 'https://platform.acme.com').
        api_base_url: REST API base (e.g., 'https://api.acme.com').
        websocket_url: WebSocket gateway base (e.g., 'wss://ws.acme.com/prod').
            The token will be appended as ?Authorization=<token>.
        platform_token: JWT token from the platform's SSO authentication.
            This is stored in memory only.
        heartbeat_interval: Seconds between heartbeats during execution (default 25).

    Returns:
        JSON confirming the connection.
    """
    tokens = _get_token_manager()

    # Store the platform token (memory only)
    tokens.store(
        token_type="platform_jwt",
        value=platform_token,
        metadata={"platform_id": platform_id},
    )

    config = PlatformConfig(
        platform_id=platform_id,
        name=name,
        base_url=base_url,
        api_base_url=api_base_url,
        websocket_url=websocket_url,
        heartbeat_interval_seconds=heartbeat_interval,
    )

    connection = PlatformConnection(
        config=config,
        is_authenticated=True,
    )

    executor = WorkflowExecutor(config, tokens)

    _platforms[platform_id] = connection
    _executors[platform_id] = executor

    logger.info("Platform connected: %s (%s)", name, platform_id)

    return json.dumps(
        {
            "platform_id": platform_id,
            "name": name,
            "status": "connected",
            "api_base_url": api_base_url,
            "hint": f"Connected. Use list_workflows('{platform_id}') to browse chains.",
        },
        indent=2,
    )


async def list_workflows(
    platform_id: Annotated[str, Field(description="The platform connection ID from connect_platform.")],
    search: Annotated[str | None, Field(description="Filter workflows by title keyword.")] = None,
    owned_by_me: Annotated[bool, Field(description="Only show workflows created by the current user.")] = False,
) -> str:
    """List available workflow chains on a connected platform.

    Browse or search for workflows that can be executed.

    Args:
        platform_id: The platform connection ID from connect_platform.
        search: Filter workflows by title keyword.
        owned_by_me: Only show workflows created by the current user.

    Returns:
        JSON with workflow summaries including IDs, titles, and descriptions.
    """
    executor = _get_executor(platform_id)

    try:
        response = await executor.list_workflows(
            search=search,
            owned_by_me=owned_by_me,
        )

        return json.dumps(
            {
                "platform_id": platform_id,
                "workflows": [
                    {
                        "workflow_id": w.workflow_id,
                        "title": w.title,
                        "description": w.description,
                        "owner": w.owner,
                        "is_owned_by_me": w.is_owned_by_me,
                    }
                    for w in response.workflows
                ],
                "total_count": response.total_count,
                "hint": "Use get_workflow_details(platform_id, workflow_id) "
                "to see components and input variables.",
            },
            indent=2,
        )

    except Exception as e:
        return json.dumps({"error": str(e), "platform_id": platform_id})


async def get_workflow_details(
    platform_id: Annotated[str, Field(description="The platform connection ID.")],
    workflow_id: Annotated[str, Field(description="UUID of the workflow to inspect.")],
) -> str:
    """Get the full definition of a workflow chain.

    Returns the components (LLMs, MCP servers, data connectors),
    input variables (template fields the user must fill), and
    the system prompt structure.

    Args:
        platform_id: The platform connection ID.
        workflow_id: UUID of the workflow to inspect.

    Returns:
        JSON with full workflow definition.
    """
    executor = _get_executor(platform_id)

    try:
        definition = await executor.get_workflow(workflow_id)

        return json.dumps(
            {
                "platform_id": platform_id,
                "workflow_id": definition.workflow_id,
                "title": definition.title,
                "description": definition.description,
                "components": [
                    {
                        "id": c.component_id,
                        "name": c.component_name,
                        "type": c.component_type,
                        "family": c.component_family,
                        "requires_config": [
                            {"key": p.key, "label": p.label, "type": p.type, "required": p.required}
                            for p in c.model_params
                        ],
                    }
                    for c in definition.components
                ],
                "input_variables": [
                    {
                        "key": v.key,
                        "label": v.label,
                        "type": v.type,
                        "required": v.required,
                    }
                    for v in definition.input_variables
                ],
                "hint": "Use execute_workflow with the workflow_id, query, "
                "modelparams, and input_variables to run this chain.",
            },
            indent=2,
        )

    except Exception as e:
        return json.dumps({"error": str(e), "platform_id": platform_id})


async def execute_workflow(
    platform_id: Annotated[str, Field(description="The platform connection ID.")],
    workflow_id: Annotated[str, Field(description="UUID of the workflow to execute.")],
    query: Annotated[str, Field(description="The user's prompt / instruction for the agent.")],
    modelparams: Annotated[
        str | None,
        Field(
            description="Optional JSON string with per-component config. Example: '{\"tool_abc\": {\"azure_devops_token\": \"...\", \"azure_devops_teams\": \"team1,team2\"}}'"
        ),
    ] = None,
    input_variables: Annotated[
        str | None,
        Field(
            description="Optional JSON string with template variable values. Example: '{\"manual_entry_123_abc\": \"value1\", \"manual_entry_123_def\": \"value2\"}'"
        ),
    ] = None,
    conversation_id: Annotated[
        str | None,
        Field(description="Optional ID to continue a previous conversation."),
    ] = None,
) -> str:
    """Execute a workflow chain and return the complete result.

    Opens a WebSocket connection, sends the execution request, collects
    all streaming responses (answer text, tool calls, stats), and returns
    the final result. The connection is closed automatically after completion.

    Args:
        platform_id: The platform connection ID.
        workflow_id: UUID of the workflow to execute.
        query: The user's prompt / instruction for the agent.
        modelparams: Optional JSON string with per-component config.
            Example: '{"tool_abc": {"azure_devops_token": "...", "azure_devops_teams": "team1,team2"}}'
        input_variables: Optional JSON string with template variable values.
            Example: '{"manual_entry_123_abc": "value1", "manual_entry_123_def": "value2"}'
        conversation_id: Optional ID to continue a previous conversation.

    Returns:
        JSON with: answer (full text), tool_calls (requests and responses),
        execution_stats, duration_seconds, and is_complete flag.
    """
    executor = _get_executor(platform_id)

    # Parse optional JSON params
    mp: dict[str, dict[str, str]] = {}
    if modelparams:
        try:
            mp = json.loads(modelparams) if isinstance(modelparams, str) else modelparams
        except json.JSONDecodeError:
            return json.dumps({"error": f"Invalid modelparams JSON: {modelparams}"})

    iv: dict[str, str] = {}
    if input_variables:
        try:
            iv = json.loads(input_variables) if isinstance(input_variables, str) else input_variables
        except json.JSONDecodeError:
            return json.dumps({"error": f"Invalid input_variables JSON: {input_variables}"})

    try:
        result = await executor.execute(
            workflow_id=workflow_id,
            query=query,
            modelparams=mp,
            input_variables=iv,
            conversation_id=conversation_id,
        )

        return json.dumps(
            {
                "platform_id": platform_id,
                "workflow_id": workflow_id,
                **result,
            },
            indent=2,
            default=str,
        )

    except Exception as e:
        logger.error("Workflow execution failed: %s", e)
        return json.dumps(
            {
                "error": str(e),
                "platform_id": platform_id,
                "workflow_id": workflow_id,
            }
        )


async def get_platform_status(
    platform_id: Annotated[str | None, Field(description="Optional specific platform to check.")] = None,
) -> str:
    """Get the status of connected enterprise platforms.

    Without platform_id, returns all connected platforms.
    With platform_id, returns detailed status for that platform.

    Args:
        platform_id: Optional specific platform to check.

    Returns:
        JSON with connection status and token validity.
    """
    tokens = _get_token_manager()

    if platform_id:
        conn = _platforms.get(platform_id)
        if not conn:
            return json.dumps({"error": f"Platform '{platform_id}' not found."})

        token_summary = tokens.summary()
        return json.dumps(
            {
                "platform_id": platform_id,
                "name": conn.config.name,
                "is_authenticated": conn.is_authenticated,
                "tokens": token_summary,
                "api_base_url": conn.config.api_base_url,
            },
            indent=2,
        )

    # All platforms
    result = []
    for pid, conn in _platforms.items():
        result.append(
            {
                "platform_id": pid,
                "name": conn.config.name,
                "is_authenticated": conn.is_authenticated,
            }
        )

    return json.dumps(
        {
            "connected_platforms": result,
            "count": len(result),
        },
        indent=2,
    )


async def disconnect_platform(
    platform_id: Annotated[str, Field(description="The platform connection ID to disconnect.")],
) -> str:
    """Disconnect from an enterprise platform and clear all tokens.

    Closes the HTTP client, clears all tokens from memory, and
    removes the platform connection.

    Args:
        platform_id: The platform connection ID to disconnect.

    Returns:
        JSON confirming disconnection.
    """
    tokens = _get_token_manager()

    # Close executor
    executor = _executors.pop(platform_id, None)
    if executor:
        await executor.close()

    # Remove connection
    _platforms.pop(platform_id, None)

    # Clear all tokens (safe even if already cleared)
    tokens.clear_all()

    logger.info("Platform disconnected: %s", platform_id)

    return json.dumps(
        {
            "platform_id": platform_id,
            "status": "disconnected",
            "message": "Connection closed and all tokens cleared from memory.",
        },
        indent=2,
    )
