#!/usr/bin/env python3
"""
Zombie Agent — Generic self-configuring ephemeral agent.

Reads ALL configuration from environment variables via config.py.
Operates in a 3-phase DAG:
  gather_context → analyze_and_execute → synthesize_and_report

Phase 1 (gather_context): Deterministic. Pull artifacts, prepare context.
Phase 2 (analyze_and_execute): LLM autonomous. Agent with tools decides.
Phase 3 (synthesize_and_report): Deterministic. Package results, upload.
"""

import json
import os

from abi_core.common.utils import abi_logging
from abi_agents.zombie.agent.config import config

# Tools loaded once at init, shared with DAG functions
ZOMBIE_TOOLS = []

# ── Idle mode ──────────────────────────────────────────────────

if config.ZOMBIE_MODE == "idle":
    import time
    abi_logging(f"[🧟] Zombie '{config.AGENT_NAME}' in IDLE mode")
    while True:
        time.sleep(60)

# ── Active mode ────────────────────────────────────────────────

abi_logging(f"[🧟→🤖] Zombie activating as '{config.AGENT_NAME}'")
abi_logging(f"[🔧] Tools: {config.TOOL_NAMES}")
abi_logging(f"[🔌] Port: {config.AGENT_PORT}")

from abi_core.agent.agent import AbiAgent
from abi_core.common import prompts


class ZombieAgent(AbiAgent):
    """Self-configuring ephemeral agent with 3-phase DAG."""

    def __init__(self):
        global ZOMBIE_TOOLS
        import asyncio
        from abi_core.common.context_loader import load_agent_context

        loop = asyncio.new_event_loop()
        ctx = loop.run_until_complete(load_agent_context(
            mcp_tool_names=config.TOOL_NAMES if config.TOOL_NAMES else None,
            library_tool_names=["write_file", "read_file", "run_shell", "list_files"],
            workspace=config.WORKSPACE,
        ))
        loop.close()

        ZOMBIE_TOOLS = ctx["lib_tools"] + ctx["mcp_tools"]

        full_prompt = prompts.ZOMBIE_AGENT_PROMPT.format(
            system_prompt=config.SYSTEM_PROMPT
        )

        abi_logging(f"System prompt injected: {config.SYSTEM_PROMPT}")

        super().__init__(
            agent_name=config.AGENT_NAME,
            description=config.AGENT_DESCRIPTION,
            llm_config=config.LLM_CONFIG,
            tools=ZOMBIE_TOOLS,
            system_prompt=full_prompt,
        )


# ── DAG: 3-phase execution ─────────────────────────────────────

from abi_core.common.tool_graph import ToolExecutionGraph, ToolGraphNode


def build_zombie_dag():
    """Build the 3-phase DAG for the zombie agent."""
    graph = ToolExecutionGraph(graph_id=f"zombie_{config.AGENT_NAME}")

    graph.add_node(ToolGraphNode(
        id="gather_context", fn=gather_context,
        input_map={"query": "$input.query"}, output_key="gather_context",
    ))
    graph.add_node(ToolGraphNode(
        id="analyze_and_execute", fn=analyze_and_execute,
        input_map={"context": "$gather_context", "query": "$input.query"},
        output_key="analyze_and_execute", depends_on=["gather_context"],
    ))
    graph.add_node(ToolGraphNode(
        id="synthesize_and_report", fn=synthesize_and_report,
        input_map={"execution_result": "$analyze_and_execute", "context": "$gather_context"},
        output_key="synthesize_and_report", depends_on=["analyze_and_execute"],
    ))
    return graph


async def gather_context(query):
    """Phase 1: Pull specific artifacts from MinIO, prepare workspace."""
    from abi_core.common.context_loader import load_agent_context

    ctx = await load_agent_context(
        artifact_keys=config.ARTIFACT_KEYS if config.ARTIFACT_KEYS else None,
        workspace=config.WORKSPACE,
    )
    return {
        "query": query,
        "workspace": config.WORKSPACE,
        "artifacts": ctx["artifacts"],
        "tools_available": config.TOOL_NAMES,
        "system_prompt": config.SYSTEM_PROMPT,
    }


async def analyze_and_execute(context, query):
    """Phase 2: LLM autonomous execution with tools."""
    from abi_core.agent.llm_provider import invoke
    from abi_core.common.context_loader import build_execution_prompt

    execution_prompt = build_execution_prompt(
        query=query,
        tools=ZOMBIE_TOOLS,
        artifacts=context.get("artifacts", []),
        workspace=config.WORKSPACE,
    )
    result = await invoke(
        config.LLM_CONFIG, execution_prompt,
        tools=ZOMBIE_TOOLS, system_prompt=config.SYSTEM_PROMPT,
    )
    return {"result": result, "query": query}


async def synthesize_and_report(execution_result, context):
    """Phase 3: Upload generated artifacts and flush logs."""
    from abi_core.common.artifact_store import upload_workspace_artifacts

    uploaded_artifacts = await upload_workspace_artifacts(
        agent_name=config.AGENT_NAME,
        workspace=config.WORKSPACE,
        exclude=context.get("artifacts", []),
    )

    return {
        "status": "completed",
        "result": execution_result.get("result", ""),
        "uploaded_artifacts": uploaded_artifacts,
        "agent": config.AGENT_NAME,
    }


# ── Start ───────────────────────────────────────────────────────

from abi_core.agent.agent_factory import agent_factory


def run_zombie():
    """Start the zombie agent with 3-phase DAG + post-response cleanup."""
    agent_instance = ZombieAgent()
    agent_instance.tool_graph = build_zombie_dag()

    # Register post-response cleanup: deregister + destroy after A2A response is sent
    if os.getenv("EPHEMERAL_SELF_DESTROY", "true").lower() == "true":
        agent_instance.on_response_sent = _post_response_cleanup

    abi_logging("[🔧] 3-phase DAG built: gather_context → analyze_and_execute → synthesize_and_report")

    agent_card = config.build_agent_card()

    abi_logging(f"[🚀] Starting zombie '{config.AGENT_NAME}' on port {config.AGENT_PORT}")
    return agent_factory(agent_instance, config, agent_card, host="0.0.0.0")


async def _post_response_cleanup():
    """Cleanup after A2A response is sent: deregister + schedule exit."""
    from abi_core.common.context_loader import self_deregister
    await self_deregister(agent_name=config.AGENT_NAME, destroy=True)


if __name__ == "__main__":
    exit(run_zombie())
