# Copyright (c) 2025 Patrick Motsch
# Main execution engine for automation2 graphs.

import logging
from datetime import datetime, timezone
from typing import Dict, Any, List, Set, Optional

from modules.workflows.automation2.graphUtils import (
    parseGraph,
    buildConnectionMap,
    validateGraph,
    topoSort,
    getInputSources,
)

from modules.workflows.automation2.executors import (
    TriggerExecutor,
    FlowExecutor,
    DataExecutor,
    ActionNodeExecutor,
    InputExecutor,
    PauseForHumanTaskError,
    PauseForEmailWaitError,
)
from modules.features.automation2.nodeDefinitions import STATIC_NODE_TYPES

logger = logging.getLogger(__name__)


def _getNodeTypeIds(services: Any = None) -> Set[str]:
    """Collect all known node type IDs from static definitions."""
    return {n["id"] for n in STATIC_NODE_TYPES}


def _getExecutor(
    nodeType: str,
    services: Any,
    automation2_interface: Optional[Any] = None,
) -> Any:
    """Dispatch to correct executor based on node type."""
    if nodeType.startswith("trigger."):
        return TriggerExecutor()
    if nodeType.startswith("flow."):
        return FlowExecutor()
    if nodeType.startswith("data."):
        return DataExecutor()
    if nodeType.startswith("ai.") or nodeType.startswith("email.") or nodeType.startswith("sharepoint."):
        return ActionNodeExecutor(services)
    if nodeType.startswith("input.") and automation2_interface:
        return InputExecutor(automation2_interface)
    return None


async def executeGraph(
    graph: Dict[str, Any],
    services: Any,
    workflowId: str = None,
    instanceId: str = None,
    userId: str = None,
    mandateId: str = None,
    automation2_interface: Optional[Any] = None,
    initialNodeOutputs: Optional[Dict[str, Any]] = None,
    startAfterNodeId: Optional[str] = None,
    runId: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Execute automation2 graph. Returns { success, nodeOutputs, error?, stopped? }.
    When an input node is reached and automation2_interface is provided, creates a task,
    pauses the run, and returns { success: False, paused: True, taskId, runId }.
    For resume: pass initialNodeOutputs (with result for the human node) and startAfterNodeId.
    """
    logger.info(
        "executeGraph start: instanceId=%s workflowId=%s userId=%s mandateId=%s resume=%s",
        instanceId,
        workflowId,
        userId,
        mandateId,
        startAfterNodeId is not None,
    )
    nodeTypeIds = _getNodeTypeIds(services)
    logger.debug("executeGraph nodeTypeIds (%d): %s", len(nodeTypeIds), sorted(nodeTypeIds))
    errors = validateGraph(graph, nodeTypeIds)
    if errors:
        logger.warning("executeGraph validation failed: %s", errors)
        return {"success": False, "error": "; ".join(errors), "nodeOutputs": {}}

    nodes, connections = parseGraph(graph)[:2]
    connectionMap = buildConnectionMap(connections)
    inputSources = {n["id"]: getInputSources(n["id"], connectionMap) for n in nodes if n.get("id")}
    logger.info(
        "executeGraph parsed: nodes=%d connections=%d connectionMap_targets=%s",
        len(nodes),
        len(connections),
        list(connectionMap.keys()),
    )

    ordered = topoSort(nodes, connectionMap)
    ordered_ids = [n.get("id") for n in ordered if n.get("id")]
    logger.info("executeGraph topoSort order: %s", ordered_ids)

    nodeOutputs: Dict[str, Any] = dict(initialNodeOutputs or {})
    is_resume = startAfterNodeId is not None
    if not runId and automation2_interface and workflowId and not is_resume:
        run_context = {
            "connectionMap": connectionMap,
            "inputSources": inputSources,
            "orderedNodeIds": ordered_ids,
        }
        if userId:
            run_context["ownerId"] = userId
        if mandateId:
            run_context["mandateId"] = mandateId
        if instanceId:
            run_context["instanceId"] = instanceId
        run = automation2_interface.createRun(
            workflowId=workflowId,
            nodeOutputs=nodeOutputs,
            context=run_context,
        )
        runId = run.get("id") if run else None
        logger.info("executeGraph created run %s", runId)

    context = {
        "workflowId": workflowId,
        "instanceId": instanceId,
        "userId": userId,
        "mandateId": mandateId,
        "nodeOutputs": nodeOutputs,
        "connectionMap": connectionMap,
        "inputSources": inputSources,
        "services": services,
        "_runId": runId,
        "_orderedNodes": ordered,
    }

    skip_until_passed = bool(startAfterNodeId)
    for i, node in enumerate(ordered):
        if skip_until_passed:
            if node.get("id") == startAfterNodeId:
                skip_until_passed = False
            continue
        if context.get("_stopped"):
            logger.info("executeGraph stopped early (flow.stop) at step %d", i)
            break
        nodeId = node.get("id")
        nodeType = node.get("type", "")
        executor = _getExecutor(nodeType, services, automation2_interface)
        logger.info(
            "executeGraph step %d/%d: nodeId=%s nodeType=%s executor=%s",
            i + 1,
            len(ordered),
            nodeId,
            nodeType,
            type(executor).__name__ if executor else "None",
        )
        if not executor:
            nodeOutputs[nodeId] = None
            logger.debug("executeGraph node %s: no executor, output=None", nodeId)
            continue
        try:
            result = await executor.execute(node, context)
            nodeOutputs[nodeId] = result
            logger.info(
                "executeGraph node %s done: result_type=%s result_keys=%s",
                nodeId,
                type(result).__name__,
                list(result.keys()) if isinstance(result, dict) else "n/a",
            )
        except PauseForHumanTaskError as e:
            logger.info("executeGraph paused for human task %s", e.taskId)
            return {
                "success": False,
                "paused": True,
                "taskId": e.taskId,
                "runId": e.runId,
                "nodeId": e.nodeId,
                "nodeOutputs": dict(nodeOutputs),
            }
        except PauseForEmailWaitError as e:
            logger.info("executeGraph paused for email wait (run %s, node %s)", e.runId, e.nodeId)
            # Start email poller on-demand (only runs while workflows wait for email)
            try:
                from modules.interfaces.interfaceDbApp import getRootInterface
                from modules.features.automation2.emailPoller import ensureRunning
                root = getRootInterface()
                event_user = root.getUserByUsername("event") if root else None
                if event_user:
                    ensureRunning(event_user)
            except Exception as poll_err:
                logger.warning("Could not start email poller: %s", poll_err)
            paused_at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
            run_ctx = {
                "connectionMap": context.get("connectionMap"),
                "inputSources": context.get("inputSources"),
                "orderedNodeIds": [n.get("id") for n in context.get("_orderedNodes", []) if n.get("id")],
                "waitReason": "email",
                "waitConfig": e.waitConfig,
                "pausedAt": paused_at,
                "lastCheckedAt": None,
                "ownerId": context.get("userId"),
                "mandateId": context.get("mandateId"),
                "instanceId": context.get("instanceId"),
            }
            automation2_interface.updateRun(
                e.runId,
                status="paused",
                nodeOutputs=dict(nodeOutputs),
                currentNodeId=e.nodeId,
                context=run_ctx,
            )
            return {
                "success": False,
                "paused": True,
                "waitReason": "email",
                "runId": e.runId,
                "nodeId": e.nodeId,
                "nodeOutputs": dict(nodeOutputs),
            }
        except Exception as e:
            logger.exception("executeGraph node %s (%s) FAILED: %s", nodeId, nodeType, e)
            nodeOutputs[nodeId] = {"error": str(e), "success": False}
            if runId and automation2_interface:
                automation2_interface.updateRun(runId, status="failed", nodeOutputs=nodeOutputs)
            return {
                "success": False,
                "error": str(e),
                "nodeOutputs": nodeOutputs,
                "failedNode": nodeId,
            }

    if runId and automation2_interface:
        automation2_interface.updateRun(runId, status="completed", nodeOutputs=nodeOutputs)
    logger.info(
        "executeGraph complete: success=True nodeOutputs_keys=%s stopped=%s",
        list(nodeOutputs.keys()),
        context.get("_stopped", False),
    )
    return {
        "success": True,
        "nodeOutputs": nodeOutputs,
        "stopped": context.get("_stopped", False),
    }