# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
"""
Workflow Toolbox - AI-assisted graph manipulation tools for the GraphicalEditor.
Tools: readWorkflowGraph, addNode, removeNode, connectNodes, setNodeParameter,
       listAvailableNodeTypes, validateGraph, listWorkflowHistory, readWorkflowMessages.

Conventions enforced here (matches coreTools / actionToolAdapter):
  - Every ``ToolResult(...)`` provides ``toolCallId`` and ``toolName`` (pydantic
    requires both); ``ToolRegistry.dispatch`` overwrites ``toolCallId`` later
    but the model still validates at construction.
  - ``ToolResult.data`` is a ``str``; structured payloads are JSON-encoded.
  - ``workflowId`` and ``instanceId`` are auto-injected from the agent
    ``context`` dict (``workflowId``, ``featureInstanceId``) when the model
    omits them — the editor agent always runs in exactly one workflow.
"""

import json
import logging
import uuid
from typing import Dict, Any, List, Tuple

from modules.serviceCenter.services.serviceAgent.datamodelAgent import ToolResult

logger = logging.getLogger(__name__)

TOOLBOX_ID = "workflow"


def _toData(payload: Any) -> str:
    """Encode a structured payload into ToolResult.data (which is a string)."""
    if isinstance(payload, str):
        return payload
    try:
        return json.dumps(payload, default=str, ensure_ascii=False)
    except Exception:
        return str(payload)


def _err(toolName: str, message: str) -> ToolResult:
    return ToolResult(toolCallId="", toolName=toolName, success=False, error=message)


def _ok(toolName: str, payload: Any) -> ToolResult:
    return ToolResult(toolCallId="", toolName=toolName, success=True, data=_toData(payload))


def _resolveIds(params: Dict[str, Any], context: Any) -> Tuple[str, str]:
    """Return (workflowId, instanceId), auto-injecting from context when missing.

    The editor agent context (``agentLoop._executeToolCalls``) is a dict with
    ``workflowId`` and ``featureInstanceId`` — use them as defaults so the
    model doesn't have to re-state the ids on every tool call.
    """
    ctx: Dict[str, Any] = context if isinstance(context, dict) else {}
    workflowId = params.get("workflowId") or ctx.get("workflowId") or ""
    instanceId = (
        params.get("instanceId")
        or ctx.get("featureInstanceId")
        or ctx.get("instanceId")
        or ""
    )
    return workflowId, instanceId


def _resolveUser(context: Any):
    """Return the User object for the current agent context (lazy DB fetch)."""
    if not isinstance(context, dict):
        return getattr(context, "user", None)
    user = context.get("user")
    if user is not None:
        return user
    userId = context.get("userId")
    if not userId:
        return None
    try:
        from modules.interfaces.interfaceDbApp import getRootInterface
        return getRootInterface().getUser(str(userId))
    except Exception as e:
        logger.warning("workflowTools: could not resolve user %s: %s", userId, e)
        return None


def _resolveMandateId(context: Any) -> str:
    if not isinstance(context, dict):
        return getattr(context, "mandateId", "") or ""
    return context.get("mandateId") or ""


def _getInterface(context: Any, instanceId: str):
    from modules.features.graphicalEditor.interfaceFeatureGraphicalEditor import getGraphicalEditorInterface
    return getGraphicalEditorInterface(_resolveUser(context), _resolveMandateId(context), instanceId)


async def _readWorkflowGraph(params: Dict[str, Any], context: Any) -> ToolResult:
    """Read the current workflow graph (nodes and connections)."""
    name = "readWorkflowGraph"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        if not workflowId or not instanceId:
            return _err(name, "workflowId and instanceId required (and not present in agent context)")

        iface = _getInterface(context, instanceId)
        wf = iface.getWorkflow(workflowId)
        if not wf:
            return _err(name, f"Workflow {workflowId} not found")

        graph = wf.get("graph", {}) or {}
        nodes = graph.get("nodes", []) or []
        connections = graph.get("connections", []) or []
        return _ok(name, {
            "workflowId": workflowId,
            "label": wf.get("label", ""),
            "nodeCount": len(nodes),
            "connectionCount": len(connections),
            "nodes": [
                {"id": n.get("id"), "type": n.get("type"), "title": n.get("title", "")}
                for n in nodes
            ],
            "connections": connections,
        })
    except Exception as e:
        logger.exception("readWorkflowGraph failed: %s", e)
        return _err(name, str(e))


async def _addNode(params: Dict[str, Any], context: Any) -> ToolResult:
    """Add a node to the workflow graph."""
    name = "addNode"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        nodeType = params.get("nodeType")
        if not workflowId or not instanceId or not nodeType:
            return _err(name, "workflowId, instanceId, and nodeType required")

        iface = _getInterface(context, instanceId)
        wf = iface.getWorkflow(workflowId)
        if not wf:
            return _err(name, f"Workflow {workflowId} not found")

        graph = dict(wf.get("graph", {}) or {})
        nodes = list(graph.get("nodes", []) or [])

        nodeId = params.get("nodeId") or str(uuid.uuid4())[:8]
        title = params.get("title", "")
        nodeParams = params.get("parameters", {}) or {}
        position = params.get("position") or {"x": len(nodes) * 200, "y": 100}

        newNode = {
            "id": nodeId,
            "type": nodeType,
            "title": title,
            "parameters": nodeParams,
            "position": position,
        }
        nodes.append(newNode)
        graph["nodes"] = nodes

        iface.updateWorkflow(workflowId, {"graph": graph})
        return _ok(name, {
            "nodeId": nodeId,
            "nodeType": nodeType,
            "message": f"Node '{title or nodeType}' added",
        })
    except Exception as e:
        logger.exception("addNode failed: %s", e)
        return _err(name, str(e))


async def _removeNode(params: Dict[str, Any], context: Any) -> ToolResult:
    """Remove a node and its connections from the workflow graph."""
    name = "removeNode"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        nodeId = params.get("nodeId")
        if not workflowId or not instanceId or not nodeId:
            return _err(name, "workflowId, instanceId, and nodeId required")

        iface = _getInterface(context, instanceId)
        wf = iface.getWorkflow(workflowId)
        if not wf:
            return _err(name, f"Workflow {workflowId} not found")

        graph = dict(wf.get("graph", {}) or {})
        nodes = [n for n in (graph.get("nodes", []) or []) if n.get("id") != nodeId]
        connections = [
            c for c in (graph.get("connections", []) or [])
            if c.get("source") != nodeId and c.get("target") != nodeId
        ]
        graph["nodes"] = nodes
        graph["connections"] = connections

        iface.updateWorkflow(workflowId, {"graph": graph})
        return _ok(name, {"nodeId": nodeId, "message": f"Node {nodeId} removed"})
    except Exception as e:
        logger.exception("removeNode failed: %s", e)
        return _err(name, str(e))


async def _connectNodes(params: Dict[str, Any], context: Any) -> ToolResult:
    """Connect two nodes in the workflow graph."""
    name = "connectNodes"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        sourceId = params.get("sourceId")
        targetId = params.get("targetId")
        if not workflowId or not instanceId or not sourceId or not targetId:
            return _err(name, "workflowId, instanceId, sourceId, and targetId required")

        iface = _getInterface(context, instanceId)
        wf = iface.getWorkflow(workflowId)
        if not wf:
            return _err(name, f"Workflow {workflowId} not found")

        graph = dict(wf.get("graph", {}) or {})
        connections = list(graph.get("connections", []) or [])
        newConn = {
            "source": sourceId,
            "target": targetId,
            "sourceOutput": params.get("sourceOutput", 0),
            "targetInput": params.get("targetInput", 0),
        }
        connections.append(newConn)
        graph["connections"] = connections

        iface.updateWorkflow(workflowId, {"graph": graph})
        return _ok(name, {"connection": newConn, "message": f"Connected {sourceId} -> {targetId}"})
    except Exception as e:
        logger.exception("connectNodes failed: %s", e)
        return _err(name, str(e))


async def _setNodeParameter(params: Dict[str, Any], context: Any) -> ToolResult:
    """Set a parameter on a node."""
    name = "setNodeParameter"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        nodeId = params.get("nodeId")
        paramName = params.get("parameterName")
        paramValue = params.get("parameterValue")
        if not workflowId or not instanceId or not nodeId or not paramName:
            return _err(name, "workflowId, instanceId, nodeId, and parameterName required")

        iface = _getInterface(context, instanceId)
        wf = iface.getWorkflow(workflowId)
        if not wf:
            return _err(name, f"Workflow {workflowId} not found")

        graph = dict(wf.get("graph", {}) or {})
        nodes = list(graph.get("nodes", []) or [])
        found = False
        for n in nodes:
            if n.get("id") == nodeId:
                nodeParams = dict(n.get("parameters", {}) or {})
                nodeParams[paramName] = paramValue
                n["parameters"] = nodeParams
                found = True
                break

        if not found:
            return _err(name, f"Node {nodeId} not found in graph")

        graph["nodes"] = nodes
        iface.updateWorkflow(workflowId, {"graph": graph})
        return _ok(name, {
            "nodeId": nodeId,
            "parameter": paramName,
            "message": f"Parameter '{paramName}' set",
        })
    except Exception as e:
        logger.exception("setNodeParameter failed: %s", e)
        return _err(name, str(e))


def _coerceLabel(rawLabel: Any, fallback: str) -> str:
    """Normalize a node label which may be a string, dict {locale: str}, or other."""
    if isinstance(rawLabel, str):
        return rawLabel
    if isinstance(rawLabel, dict):
        for key in ("en", "de", "fr"):
            value = rawLabel.get(key)
            if isinstance(value, str) and value:
                return value
        for value in rawLabel.values():
            if isinstance(value, str) and value:
                return value
    return fallback


async def _listAvailableNodeTypes(params: Dict[str, Any], context: Any) -> ToolResult:
    """List all available node types for the flow builder."""
    name = "listAvailableNodeTypes"
    try:
        from modules.features.graphicalEditor.nodeDefinitions import STATIC_NODE_TYPES
        nodeTypes = []
        for n in STATIC_NODE_TYPES:
            if not isinstance(n, dict):
                continue
            nodeId = n.get("id") or ""
            nodeTypes.append({
                "id": nodeId,
                "category": n.get("category"),
                "label": _coerceLabel(n.get("label"), nodeId),
            })
        return _ok(name, {"nodeTypes": nodeTypes, "count": len(nodeTypes)})
    except Exception as e:
        logger.exception("listAvailableNodeTypes failed: %s", e)
        return _err(name, str(e))


async def _validateGraph(params: Dict[str, Any], context: Any) -> ToolResult:
    """Validate a workflow graph for common issues."""
    name = "validateGraph"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        if not workflowId or not instanceId:
            return _err(name, "workflowId and instanceId required")

        iface = _getInterface(context, instanceId)
        wf = iface.getWorkflow(workflowId)
        if not wf:
            return _err(name, f"Workflow {workflowId} not found")

        graph = wf.get("graph", {}) or {}
        nodes = graph.get("nodes", []) or []
        connections = graph.get("connections", []) or []
        issues: List[str] = []

        nodeIds = {n.get("id") for n in nodes}
        if not nodes:
            issues.append("Graph has no nodes")

        hasTrigger = any((n.get("type") or "").startswith("trigger.") for n in nodes)
        if not hasTrigger:
            issues.append("No trigger node found")

        for c in connections:
            if c.get("source") not in nodeIds:
                issues.append(f"Connection source '{c.get('source')}' not found")
            if c.get("target") not in nodeIds:
                issues.append(f"Connection target '{c.get('target')}' not found")

        connectedNodes = set()
        for c in connections:
            connectedNodes.add(c.get("source"))
            connectedNodes.add(c.get("target"))
        orphans = [
            n.get("id") for n in nodes
            if n.get("id") not in connectedNodes and not (n.get("type") or "").startswith("trigger.")
        ]
        if orphans:
            issues.append(f"Orphan nodes (not connected): {', '.join(orphans)}")

        return _ok(name, {
            "valid": len(issues) == 0,
            "issues": issues,
            "nodeCount": len(nodes),
            "connectionCount": len(connections),
        })
    except Exception as e:
        logger.exception("validateGraph failed: %s", e)
        return _err(name, str(e))


async def _listWorkflowHistory(params: Dict[str, Any], context: Any) -> ToolResult:
    """List versions (history) for a workflow."""
    name = "listWorkflowHistory"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        if not workflowId or not instanceId:
            return _err(name, "workflowId and instanceId required")
        iface = _getInterface(context, instanceId)
        versions = iface.getVersions(workflowId) or []
        return _ok(name, {
            "workflowId": workflowId,
            "versions": [
                {
                    "id": v.get("id"),
                    "versionNumber": v.get("versionNumber"),
                    "status": v.get("status"),
                    "publishedAt": v.get("publishedAt"),
                    "publishedBy": v.get("publishedBy"),
                }
                for v in versions
            ],
        })
    except Exception as e:
        logger.exception("listWorkflowHistory failed: %s", e)
        return _err(name, str(e))


async def _readWorkflowMessages(params: Dict[str, Any], context: Any) -> ToolResult:
    """Read recent run logs/messages for a workflow."""
    name = "readWorkflowMessages"
    try:
        workflowId, instanceId = _resolveIds(params, context)
        if not workflowId or not instanceId:
            return _err(name, "workflowId and instanceId required")
        iface = _getInterface(context, instanceId)
        from modules.features.graphicalEditor.datamodelFeatureGraphicalEditor import AutoRun
        runs = iface.db.getRecordset(AutoRun, recordFilter={"workflowId": workflowId}) or []
        runSummaries = []
        for r in sorted(runs, key=lambda x: x.get("startedAt") or 0, reverse=True)[:10]:
            runSummaries.append({
                "runId": r.get("id"),
                "status": r.get("status"),
                "startedAt": r.get("startedAt"),
                "completedAt": r.get("completedAt"),
                "error": r.get("error"),
            })
        return _ok(name, {"workflowId": workflowId, "recentRuns": runSummaries})
    except Exception as e:
        logger.exception("readWorkflowMessages failed: %s", e)
        return _err(name, str(e))


def getWorkflowToolDefinitions() -> List[Dict[str, Any]]:
    """Return tool definitions for registration in the ToolRegistry.

    Note: ``workflowId`` and ``instanceId`` are NOT marked ``required`` —
    they are auto-injected from the agent context by ``_resolveIds``. The
    model may still pass them explicitly (e.g. to target a different
    workflow) but doesn't have to repeat them on every call.
    """
    _idFields = {
        "workflowId": {"type": "string", "description": "Workflow ID (defaults to the current editor workflow)"},
        "instanceId": {"type": "string", "description": "Feature instance ID (defaults to the current editor instance)"},
    }
    return [
        {
            "name": "readWorkflowGraph",
            "handler": _readWorkflowGraph,
            "description": "Read the current workflow graph (nodes and connections). Always call this first to understand the current state before making changes.",
            "parameters": {
                "type": "object",
                "properties": {**_idFields},
                "required": [],
            },
            "readOnly": True,
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "addNode",
            "handler": _addNode,
            "description": "Add a node to the workflow graph.",
            "parameters": {
                "type": "object",
                "properties": {
                    **_idFields,
                    "nodeType": {"type": "string", "description": "Node type id (e.g. ai.chat, email.send) — use listAvailableNodeTypes to discover"},
                    "title": {"type": "string", "description": "Human-readable title"},
                    "parameters": {"type": "object", "description": "Node parameters"},
                    "position": {"type": "object", "description": "Canvas position {x, y}"},
                    "nodeId": {"type": "string", "description": "Optional explicit node id"},
                },
                "required": ["nodeType"],
            },
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "removeNode",
            "handler": _removeNode,
            "description": "Remove a node and its connections from the graph.",
            "parameters": {
                "type": "object",
                "properties": {
                    **_idFields,
                    "nodeId": {"type": "string", "description": "ID of the node to remove"},
                },
                "required": ["nodeId"],
            },
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "connectNodes",
            "handler": _connectNodes,
            "description": "Connect two nodes in the graph (source -> target).",
            "parameters": {
                "type": "object",
                "properties": {
                    **_idFields,
                    "sourceId": {"type": "string"},
                    "targetId": {"type": "string"},
                    "sourceOutput": {"type": "integer", "default": 0},
                    "targetInput": {"type": "integer", "default": 0},
                },
                "required": ["sourceId", "targetId"],
            },
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "setNodeParameter",
            "handler": _setNodeParameter,
            "description": "Set a single parameter on a node.",
            "parameters": {
                "type": "object",
                "properties": {
                    **_idFields,
                    "nodeId": {"type": "string"},
                    "parameterName": {"type": "string"},
                    "parameterValue": {"description": "Value to set (any type)"},
                },
                "required": ["nodeId", "parameterName", "parameterValue"],
            },
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "listAvailableNodeTypes",
            "handler": _listAvailableNodeTypes,
            "description": "List all available node types for the flow builder. Call this once to discover ids before using addNode.",
            "parameters": {"type": "object", "properties": {}},
            "readOnly": True,
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "validateGraph",
            "handler": _validateGraph,
            "description": "Validate a workflow graph for common issues (missing trigger, dangling connections, orphans).",
            "parameters": {
                "type": "object",
                "properties": {**_idFields},
                "required": [],
            },
            "readOnly": True,
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "listWorkflowHistory",
            "handler": _listWorkflowHistory,
            "description": "List version history for a workflow (AutoVersion entries).",
            "parameters": {
                "type": "object",
                "properties": {**_idFields},
                "required": [],
            },
            "readOnly": True,
            "toolSet": TOOLBOX_ID,
        },
        {
            "name": "readWorkflowMessages",
            "handler": _readWorkflowMessages,
            "description": "Read recent run logs and status for a workflow.",
            "parameters": {
                "type": "object",
                "properties": {**_idFields},
                "required": [],
            },
            "readOnly": True,
            "toolSet": TOOLBOX_ID,
        },
    ]