"""Workflow executor -- high-level API for running enterprise workflow chains.

This module ties together the WebSocket adapter, token manager, and REST client
to provide a clean interface for:
  1. Fetching workflow definitions from the platform REST API
  2. Executing workflows via WebSocket (open-send-receive-close)
  3. Returning structured results with answer text, tool calls, and stats

Usage:
    executor = WorkflowExecutor(platform_config, token_manager)
    workflows = await executor.list_workflows(search="impact analysis")
    definition = await executor.get_workflow(workflow_id)
    result = await executor.execute(workflow_id, query, modelparams, input_vars)
"""

from __future__ import annotations

import json
import logging
from typing import Any

import httpx

from mcp_bridge.platforms.auth.token_manager import TokenManager
from mcp_bridge.platforms.models.workflow import (
    PlatformConfig,
    WorkflowDefinition,
    WorkflowListResponse,
    WorkflowSummary,
)
from mcp_bridge.platforms.websocket_adapter import (
    StreamBuffer,
    WebSocketAdapter,
    WebSocketConfig,
)

logger = logging.getLogger(__name__)


class WorkflowExecutor:
    """Orchestrates workflow discovery and execution against an enterprise platform.

    Lifecycle per execution:
      1. Validate tokens via TokenManager
      2. (Optional) Fetch workflow definition via REST to know input variables
      3. Open WebSocket → send SendMessageV2 → collect response → close
      4. Return structured result

    This class is stateless across executions -- each ``execute()`` call
    opens and closes its own WebSocket connection.
    """

    def __init__(self, config: PlatformConfig, token_manager: TokenManager) -> None:
        self._config = config
        self._tokens = token_manager
        self._http: httpx.AsyncClient | None = None

    # ------------------------------------------------------------------
    # REST API -- workflow discovery
    # ------------------------------------------------------------------

    async def _get_http_client(self) -> httpx.AsyncClient:
        """Get or create an HTTP client with auth headers."""
        platform_token = self._tokens.get("platform_jwt")
        if not platform_token:
            raise RuntimeError("Platform JWT not available. Authenticate first.")

        if self._http is None or self._http.is_closed:
            self._http = httpx.AsyncClient(
                base_url=self._config.api_base_url,
                headers={
                    "Authorization": f"Bearer {platform_token}",
                    "Content-Type": "application/json",
                    "Accept": "application/json",
                },
                timeout=30.0,
            )
        else:
            # Refresh the auth header in case the token was rotated
            self._http.headers["Authorization"] = f"Bearer {platform_token}"

        return self._http

    async def list_workflows(
        self,
        search: str | None = None,
        owned_by_me: bool = False,
        page_size: int = 20,
        sort: str = "created_on",
        order: str = "desc",
    ) -> WorkflowListResponse:
        """List available workflow chains from the platform.

        Args:
            search: Filter by title keyword.
            owned_by_me: Only show workflows created by the current user.
            page_size: Number of results per page.
            sort: Sort field (created_on, updated_on, title).
            order: Sort order (asc, desc).

        Returns:
            WorkflowListResponse with workflow summaries.
        """
        client = await self._get_http_client()

        params: dict[str, Any] = {
            "page_size": page_size,
            "sort": sort,
            "order": order,
        }
        if search:
            params["title"] = search
        if owned_by_me:
            params["is_owned_by_me"] = "true"

        resp = await client.get("/v2/workflow_trimmed", params=params)
        resp.raise_for_status()
        data = resp.json()

        # The API returns a list directly or a paginated wrapper
        workflows_raw = data if isinstance(data, list) else data.get("workflows", data.get("items", []))

        workflows = []
        for item in workflows_raw:
            workflows.append(
                WorkflowSummary(
                    workflow_id=item.get("workflow_id", item.get("id", "")),
                    title=item.get("title", ""),
                    description=item.get("description", ""),
                    owner=item.get("owner", item.get("created_by", "")),
                    is_owned_by_me=item.get("is_owned_by_me", False),
                    component_count=len(item.get("components", [])),
                )
            )

        return WorkflowListResponse(
            workflows=workflows,
            total_count=len(workflows),
            page_size=page_size,
            has_more=len(workflows) >= page_size,
        )

    async def get_workflow(self, workflow_id: str) -> WorkflowDefinition:
        """Fetch the full definition of a workflow chain.

        Args:
            workflow_id: UUID of the workflow.

        Returns:
            WorkflowDefinition with components, input variables, and system prompt.
        """
        client = await self._get_http_client()

        resp = await client.get(f"/v2/workflow/{workflow_id}")
        resp.raise_for_status()
        data = resp.json()

        return self._parse_workflow_definition(data)

    # ------------------------------------------------------------------
    # WebSocket -- workflow execution
    # ------------------------------------------------------------------

    async def execute(
        self,
        workflow_id: str,
        query: str,
        modelparams: dict[str, dict[str, str]] | None = None,
        input_variables: dict[str, str] | None = None,
        conversation_id: str | None = None,
        is_persistence_allowed: bool = False,
    ) -> dict[str, Any]:
        """Execute a workflow chain and return the complete result.

        Opens a WebSocket, sends the execution payload, collects all
        streamed messages, and closes the connection.

        Args:
            workflow_id: UUID of the workflow to run.
            query: User prompt for the agent.
            modelparams: Per-component tokens and config (e.g., ADO JWT, GitHub PAT).
            input_variables: Template variable values.
            conversation_id: Continue a previous conversation.
            is_persistence_allowed: Whether the platform should persist history.

        Returns:
            Dict with: answer, tool_calls, execution_stats, reasoning,
            duration_seconds, is_complete.
        """
        platform_token = self._tokens.get("platform_jwt")
        if not platform_token:
            raise RuntimeError("Platform JWT not available. Authenticate first.")

        ws_config = WebSocketConfig(
            url_template=self._config.websocket_url + "?Authorization={token}",
            heartbeat_interval=self._config.heartbeat_interval_seconds,
        )
        adapter = WebSocketAdapter(ws_config)

        logger.info(
            "Executing workflow: id=%s, query_len=%d",
            workflow_id,
            len(query),
        )

        result = await adapter.run(
            token=platform_token,
            workflow_id=workflow_id,
            query=query,
            modelparams=modelparams,
            input_variables=input_variables,
            conversation_id=conversation_id,
            is_persistence_allowed=is_persistence_allowed,
        )

        logger.info(
            "Execution result: complete=%s, duration=%.1fs, tool_calls=%d",
            result.get("is_complete"),
            result.get("duration_seconds", 0),
            result.get("tool_calls", {}).get("total", 0),
        )

        return result

    # ------------------------------------------------------------------
    # Component catalog
    # ------------------------------------------------------------------

    async def list_components(self, component_type: str = "all") -> list[dict[str, Any]]:
        """Fetch the platform's component catalog.

        Args:
            component_type: Filter by type ('all', 'llm', 'tool', 'vector_db').

        Returns:
            List of component dicts with id, name, type, family.
        """
        client = await self._get_http_client()

        resp = await client.get("/v2/components", params={"component_type": component_type})
        resp.raise_for_status()
        data = resp.json()

        components = data if isinstance(data, list) else data.get("components", [])

        return [
            {
                "component_id": c.get("component_id", c.get("id", "")),
                "name": c.get("component_name", c.get("name", "")),
                "type": c.get("component_type", ""),
                "family": c.get("component_family", ""),
                "api_type": c.get("api_type", ""),
                "description": c.get("description", ""),
            }
            for c in components
        ]

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------

    async def close(self) -> None:
        """Close the HTTP client."""
        if self._http and not self._http.is_closed:
            await self._http.aclose()
            self._http = None

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _parse_workflow_definition(data: dict[str, Any]) -> WorkflowDefinition:
        """Parse a raw API response into a WorkflowDefinition model."""
        from mcp_bridge.platforms.models.workflow import (
            ComponentParam,
            ToolConfig,
            WorkflowComponent,
            WorkflowInputVariable,
        )

        components = []
        for idx, comp in enumerate(data.get("components", [])):
            params = []
            for p in comp.get("model_params", []):
                params.append(
                    ComponentParam(
                        key=p.get("key", ""),
                        type=p.get("type", "text"),
                        label=p.get("label", ""),
                        description=p.get("description", ""),
                        required=p.get("required", False),
                    )
                )

            tool_cfg = None
            raw_tool_config = comp.get("tool_config")
            if raw_tool_config and isinstance(raw_tool_config, dict):
                tool_cfg = ToolConfig(
                    endpoint=raw_tool_config.get("endpoint", ""),
                    auth_type=raw_tool_config.get("auth_type", "bearer"),
                    headers=raw_tool_config.get("headers", {}),
                )

            components.append(
                WorkflowComponent(
                    component_id=comp.get("component_id", comp.get("id", "")),
                    component_name=comp.get("component_name", comp.get("name", "")),
                    component_type=comp.get("component_type", "custom"),
                    component_family=comp.get("component_family", ""),
                    api_type=comp.get("api_type"),
                    position=comp.get("position", idx),
                    tool_config=tool_cfg,
                    model_params=params,
                )
            )

        input_vars = []
        for var in data.get("input_variables", []):
            input_vars.append(
                WorkflowInputVariable(
                    key=var.get("key", ""),
                    label=var.get("label", ""),
                    type=var.get("type", "text"),
                    description=var.get("description", ""),
                    required=var.get("required", True),
                )
            )

        return WorkflowDefinition(
            workflow_id=data.get("workflow_id", data.get("id", "")),
            title=data.get("title", ""),
            description=data.get("description", ""),
            system_prompt=data.get("system_prompt", ""),
            components=components,
            input_variables=input_vars,
            owner=data.get("owner", data.get("created_by", "")),
            is_owned_by_me=data.get("is_owned_by_me", False),
        )
