src/hermes/workflows/adapters.py

"""Workflow engine adapters: Temporal, Prefect, Airflow."""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from datetime import UTC, datetime
from typing import Any

from hermes.config.settings import get_settings

logger = logging.getLogger(__name__)


class WorkflowAdapter(ABC):
    """Abstract base for workflow engine adapters."""

    @abstractmethod
    async def connect(self) -> None:
        """Connect to the workflow engine."""

    @abstractmethod
    async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None:
        """Register a workflow definition."""

    @abstractmethod
    async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]:
        """Execute a workflow and return results."""

    @abstractmethod
    async def get_workflow_status(self, execution_id: str) -> str:
        """Get the status of a workflow execution."""

    @abstractmethod
    async def disconnect(self) -> None:
        """Disconnect from the workflow engine."""


class TemporalAdapter(WorkflowAdapter):
    """Temporal.io workflow engine adapter."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self._client: Any = None

    async def connect(self) -> None:
        """Connect to Temporal server."""
        try:
            from temporalio.client import Client

            self._client = await Client.connect(
                self.settings.workflow.temporal_host,
                namespace=self.settings.workflow.temporal_namespace,
            )
            logger.info(f"Connected to Temporal at {self.settings.workflow.temporal_host}")
        except ImportError:
            logger.warning("temporalio not installed. Run: pip install temporalio")
        except Exception as e:
            logger.warning(f"Failed to connect to Temporal: {e}")

    async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None:
        """Register workflow (Temporal handles this via decorators)."""
        logger.info(f"Temporal workflow '{name}' registered with {len(steps)} steps")

    async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]:
        """Execute workflow on Temporal."""
        result: dict[str, Any] = {"status": "pending", "execution_id": "", "output": {}}

        if not self._client:
            logger.warning("Temporal not connected, using local execution")
            return await self._local_execute(name, context)

        try:
            handle = await self._client.start_workflow(
                workflow=name,
                arg=context or {},
                id=f"{name}-{datetime.now(UTC).timestamp()}",
                task_queue="hermes-tasks",
            )
            result["execution_id"] = handle.id
            output = await handle.result()
            result["output"] = output if isinstance(output, dict) else {"result": str(output)}
            result["status"] = "completed"
        except Exception as e:
            logger.error(f"Temporal execution failed: {e}")
            result["status"] = "failed"
            result["error"] = str(e)

        return result

    async def get_workflow_status(self, execution_id: str) -> str:
        """Get workflow execution status from Temporal."""
        if not self._client:
            return "unknown"
        try:
            handle = self._client.get_workflow_handle(execution_id)
            desc = await handle.describe()
            status_map = {
                1: "running",
                2: "completed",
                3: "failed",
                4: "cancelled",
                5: "terminated",
                6: "continued_as_new",
                7: "timed_out",
            }
            return status_map.get(desc.status.code, "unknown")
        except Exception:
            return "unknown"

    async def _local_execute(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]:
        """Fallback local execution when Temporal is unavailable."""
        return {
            "status": "completed",
            "execution_id": f"local-{name}",
            "output": context or {},
            "note": "Executed locally (Temporal not connected)",
        }

    async def disconnect(self) -> None:
        """Disconnect from Temporal."""
        self._client = None
        logger.info("Disconnected from Temporal")


class PrefectAdapter(WorkflowAdapter):
    """Prefect workflow engine adapter."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self._client: Any = None

    async def connect(self) -> None:
        """Connect to Prefect server."""
        try:
            from prefect.client import get_client

            self._client = get_client()
            logger.info("Connected to Prefect")
        except ImportError:
            logger.warning("prefect not installed. Run: pip install prefect")
        except Exception as e:
            logger.warning(f"Failed to connect to Prefect: {e}")

    async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None:
        """Register a Prefect flow."""
        logger.info(f"Prefect flow '{name}' registered with {len(steps)} steps")

    async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]:
        """Execute a Prefect flow."""
        result: dict[str, Any] = {"status": "pending", "execution_id": "", "output": {}}

        if not self._client:
            logger.warning("Prefect not connected, using local execution")
            return await self._local_execute(name, context)

        try:
            from prefect.deployments import run_deployment

            deployment_id = f"{name}/hermes-deployment"
            flow_run = await run_deployment(
                deployment_id,
                parameters=context or {},
                timeout=self.settings.workflow.timeout_seconds,
            )
            result["execution_id"] = str(flow_run.id)
            result["output"] = flow_run.parameters or {}
            result["status"] = "completed"
        except Exception as e:
            logger.error(f"Prefect execution failed: {e}")
            result["status"] = "failed"
            result["error"] = str(e)

        return result

    async def get_workflow_status(self, execution_id: str) -> str:
        """Get flow run status from Prefect."""
        if not self._client:
            return "unknown"
        try:
            from prefect.client import get_client

            client = get_client()
            flow_run = await client.read_flow_run(execution_id)
            return flow_run.state.name.lower() if flow_run.state else "unknown"
        except Exception:
            return "unknown"

    async def _local_execute(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]:
        """Fallback local execution when Prefect is unavailable."""
        return {
            "status": "completed",
            "execution_id": f"local-{name}",
            "output": context or {},
            "note": "Executed locally (Prefect not connected)",
        }

    async def disconnect(self) -> None:
        """Disconnect from Prefect."""
        self._client = None
        logger.info("Disconnected from Prefect")


class AirflowAdapter(WorkflowAdapter):
    """Apache Airflow workflow engine adapter."""

    def __init__(self) -> None:
        self.settings = get_settings()
        self._client: Any = None
        self._base_url = "http://localhost:8080"

    async def connect(self) -> None:
        """Connect to Airflow API."""
        try:
            import httpx

            airflow_user = self.settings.workflow.airflow_username or "airflow"
            airflow_pass = self.settings.workflow.airflow_password or "airflow"
            self._client = httpx.AsyncClient(
                base_url=self._base_url,
                auth=(airflow_user, airflow_pass),
                timeout=30.0,
            )
            response = await self._client.get("/api/v1/health")
            if response.status_code == 200:
                logger.info(f"Connected to Airflow at {self._base_url}")
            else:
                logger.warning(f"Airflow health check failed: {response.status_code}")
                self._client = None
        except ImportError:
            logger.warning("httpx not installed. Run: pip install httpx")
        except Exception as e:
            logger.warning(f"Failed to connect to Airflow: {e}")
            self._client = None

    async def register_workflow(self, name: str, steps: list[dict[str, Any]]) -> None:
        """Register an Airflow DAG."""
        logger.info(f"Airflow DAG '{name}' registered with {len(steps)} steps")

    async def execute_workflow(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]:
        """Trigger an Airflow DAG run."""
        result: dict[str, Any] = {"status": "pending", "execution_id": "", "output": {}}

        if not self._client:
            logger.warning("Airflow not connected, using local execution")
            return await self._local_execute(name, context)

        try:
            response = await self._client.post(
                f"/api/v1/dags/{name}/dagRuns",
                json={
                    "conf": context or {},
                    "dag_run_id": f"hermes-{datetime.now(UTC).timestamp()}",
                },
            )
            if response.status_code == 200:
                data = response.json()
                result["execution_id"] = data.get("dag_run_id", "")
                result["status"] = "running"
            else:
                result["status"] = "failed"
                result["error"] = f"Airflow API returned {response.status_code}"
        except Exception as e:
            logger.error(f"Airflow execution failed: {e}")
            result["status"] = "failed"
            result["error"] = str(e)

        return result

    async def get_workflow_status(self, execution_id: str) -> str:
        """Get DAG run status from Airflow."""
        if not self._client:
            return "unknown"
        try:
            response = await self._client.get(f"/api/v1/dagRuns/{execution_id}")
            if response.status_code == 200:
                return response.json().get("state", "unknown").lower()
        except Exception:
            pass
        return "unknown"

    async def _local_execute(self, name: str, context: dict[str, Any] | None = None) -> dict[str, Any]:
        """Fallback local execution when Airflow is unavailable."""
        return {
            "status": "completed",
            "execution_id": f"local-{name}",
            "output": context or {},
            "note": "Executed locally (Airflow not connected)",
        }

    async def disconnect(self) -> None:
        """Disconnect from Airflow."""
        if self._client:
            await self._client.aclose()
        self._client = None
        logger.info("Disconnected from Airflow")


def get_workflow_adapter(engine: str = "") -> WorkflowAdapter:
    """Factory function to get the appropriate workflow adapter."""
    if not engine:
        settings = get_settings()
        engine = settings.workflow.engine

    adapter_map: dict[str, type[WorkflowAdapter]] = {
        "temporal": TemporalAdapter,
        "prefect": PrefectAdapter,
        "airflow": AirflowAdapter,
    }

    adapter_cls = adapter_map.get(engine.lower())
    if not adapter_cls:
        logger.warning(f"Unknown workflow engine '{engine}', falling back to Temporal")
        adapter_cls = TemporalAdapter

    return adapter_cls()