"""
Memory Hook System — Phase 23.5

Auto-captures agent behavior events and ingests them into the
Observation Layer → DAG Memory Graph pipeline.

Hooks:
    on_session_start(ctx)      — session begins
    on_agent_step(step)        — single reasoning step
    on_tool_use(tool_event)    — tool called + result
    on_task_complete(task)     — task finished

Usage:
    hooks = MemoryHooks(engine=memory_engine)
    hooks.on_tool_use({
        "tool": "docker",
        "input": "deploy backend",
        "output": "success",
        "task_id": "task_123",
        "scene": "deployment",
    })
"""

import logging
import time
from typing import Any, Callable, Dict, List, Optional

from openmemo.core.observation import Observation

logger = logging.getLogger("openmemo.hooks")


class HookEvent:
    """Base event wrapper passed to hook callbacks."""
    def __init__(self, event_type: str, data: dict):
        self.event_type = event_type
        self.data = data
        self.timestamp = time.time()


class MemoryHooks:
    """
    Central hook dispatcher — auto-ingests agent events into OpenMemo.

    Args:
        engine          — MemoryEngine instance (or any object with .ingest())
        agent_id        — default agent_id for all hooks
        auto_ingest     — if True, immediately ingest to engine on each hook
        on_error        — optional callback for ingest errors
    """

    def __init__(self, engine=None, agent_id: str = "",
                 auto_ingest: bool = True,
                 on_error: Optional[Callable] = None):
        self._engine = engine
        self._agent_id = agent_id
        self._auto_ingest = auto_ingest
        self._on_error = on_error
        self._listeners: Dict[str, List[Callable]] = {}
        self._ingested: List[str] = []

    def attach(self, engine) -> "MemoryHooks":
        """Attach (or replace) the memory engine."""
        self._engine = engine
        return self

    def listen(self, event_type: str, fn: Callable) -> "MemoryHooks":
        """Register a custom listener for a hook event."""
        self._listeners.setdefault(event_type, []).append(fn)
        return self

    def _emit(self, event: HookEvent):
        for fn in self._listeners.get(event.event_type, []):
            try:
                fn(event)
            except Exception as e:
                logger.warning("[hooks] listener error for %s: %s", event.event_type, e)

    def _ingest(self, obs: Observation) -> Optional[str]:
        if not self._auto_ingest or not self._engine:
            return None
        try:
            result = self._engine.ingest(obs)
            node_id = result.id if hasattr(result, "id") else str(result)
            self._ingested.append(obs.id)
            logger.debug("[hooks] ingested %s → %s", obs.id, node_id)
            return node_id
        except Exception as e:
            logger.warning("[hooks] ingest error: %s", e)
            if self._on_error:
                self._on_error(e, obs)
            return None

    # ─── Core Hooks ────────────────────────────────────────────────

    def on_session_start(self, ctx: Dict[str, Any] = None) -> Optional[str]:
        """
        Called when an agent session begins.

        Args:
            ctx: dict with optional keys: session_id, agent_id, scene, goal
        """
        ctx = ctx or {}
        agent_id = ctx.get("agent_id", self._agent_id)
        session_id = ctx.get("session_id", "")
        scene = ctx.get("scene", "")

        obs = Observation.session_event(
            event="session_start",
            session_id=session_id,
            agent_id=agent_id,
            scene=scene,
        )
        if ctx.get("goal"):
            obs.content["goal"] = ctx["goal"]

        self._emit(HookEvent("session_start", {"obs": obs, "ctx": ctx}))
        logger.info("[hooks] session_start agent=%s session=%s", agent_id, session_id)
        return self._ingest(obs)

    def on_session_end(self, ctx: Dict[str, Any] = None) -> Optional[str]:
        """Called when an agent session ends."""
        ctx = ctx or {}
        agent_id = ctx.get("agent_id", self._agent_id)
        session_id = ctx.get("session_id", "")
        scene = ctx.get("scene", "")

        obs = Observation.session_event(
            event="session_end",
            session_id=session_id,
            agent_id=agent_id,
            scene=scene,
        )
        self._emit(HookEvent("session_end", {"obs": obs, "ctx": ctx}))
        return self._ingest(obs)

    def on_agent_step(self, step: Dict[str, Any]) -> Optional[str]:
        """
        Called for each agent reasoning step.

        Args:
            step: dict with keys: step (str), reasoning (str, optional),
                  agent_id, task_id, scene
        """
        if isinstance(step, str):
            step = {"step": step}

        agent_id = step.get("agent_id", self._agent_id)
        obs = Observation.agent_step(
            step=step.get("step", ""),
            reasoning=step.get("reasoning", ""),
            agent_id=agent_id,
            task_id=step.get("task_id", ""),
            scene=step.get("scene", ""),
            confidence=step.get("confidence", 0.85),
        )

        self._emit(HookEvent("agent_step", {"obs": obs, "step": step}))
        logger.debug("[hooks] agent_step agent=%s step=%s", agent_id, step.get("step", "")[:60])
        return self._ingest(obs)

    def on_tool_use(self, tool_event: Dict[str, Any]) -> Optional[str]:
        """
        Called when an agent uses a tool.

        Args:
            tool_event: dict with keys:
                tool (str)     — tool name
                input (str)    — what was sent to the tool
                output (str)   — what the tool returned
                task_id        — associated task
                agent_id       — agent that used the tool
                scene          — memory scene
                tags           — list of tags
                confidence     — confidence score (default 0.9)
        """
        if isinstance(tool_event, str):
            tool_event = {"tool": tool_event}

        agent_id = tool_event.get("agent_id", self._agent_id)
        obs = Observation.tool_execution(
            tool=tool_event.get("tool", "unknown"),
            action=str(tool_event.get("input", tool_event.get("action", ""))),
            result=str(tool_event.get("output", tool_event.get("result", ""))),
            agent_id=agent_id,
            task_id=tool_event.get("task_id", ""),
            scene=tool_event.get("scene", ""),
            tags=tool_event.get("tags", []),
            confidence=float(tool_event.get("confidence", 0.9)),
        )

        self._emit(HookEvent("tool_use", {"obs": obs, "event": tool_event}))
        logger.debug("[hooks] tool_use agent=%s tool=%s", agent_id, tool_event.get("tool"))
        return self._ingest(obs)

    def on_task_complete(self, task: Dict[str, Any]) -> Optional[str]:
        """
        Called when a task is completed (success or failure).

        Args:
            task: dict with keys:
                task_id (str)   — task identifier
                status (str)    — "success" | "failure" | "partial"
                summary (str)   — what was accomplished
                agent_id        — agent that ran the task
                scene           — memory scene
        """
        if isinstance(task, str):
            task = {"task_id": task, "status": "complete"}

        agent_id = task.get("agent_id", self._agent_id)
        status = task.get("status", "complete")
        task_id = task.get("task_id", "")

        obs = Observation.task_event(
            event="task_complete",
            task_id=task_id,
            status=status,
            agent_id=agent_id,
            scene=task.get("scene", ""),
            confidence=0.95 if status == "success" else 0.75,
        )
        if task.get("summary"):
            obs.content["summary"] = task["summary"]

        self._emit(HookEvent("task_complete", {"obs": obs, "task": task}))
        logger.info("[hooks] task_complete agent=%s task=%s status=%s",
                    agent_id, task_id, status)
        return self._ingest(obs)

    def on_task_start(self, task: Dict[str, Any]) -> Optional[str]:
        """Called when a task begins."""
        if isinstance(task, str):
            task = {"task_id": task, "status": "started"}

        agent_id = task.get("agent_id", self._agent_id)
        obs = Observation.task_event(
            event="task_start",
            task_id=task.get("task_id", ""),
            status="started",
            agent_id=agent_id,
            scene=task.get("scene", ""),
        )
        self._emit(HookEvent("task_start", {"obs": obs, "task": task}))
        return self._ingest(obs)

    # ─── Stats ─────────────────────────────────────────────────────

    @property
    def ingested_count(self) -> int:
        return len(self._ingested)

    def status(self) -> dict:
        return {
            "agent_id": self._agent_id,
            "auto_ingest": self._auto_ingest,
            "engine_attached": self._engine is not None,
            "ingested_count": self.ingested_count,
            "listeners": {k: len(v) for k, v in self._listeners.items()},
        }
