from __future__ import annotations

import time
from typing import Any

from ..registry import get_agent
from ..run_logger import append_stage_trace
from ..schemas import StageTrace


def now_ms() -> int:
    return int(time.time() * 1000)


def append_timeline(state: dict[str, Any], event: str) -> dict[str, Any]:
    memory = dict(state.get("working_memory") or {"shared": {}, "agent_private": {}, "timeline": []})
    memory.setdefault("shared", {})
    memory.setdefault("agent_private", {})
    timeline = list(memory.get("timeline", []))
    timeline.append({"timestamp_ms": now_ms(), "event": event})
    memory["timeline"] = timeline
    return memory


def stage_index(state: dict[str, Any]) -> int:
    return len(state.get("stage_traces", [])) + 1


def record_trace(state: dict[str, Any], trace: StageTrace) -> list[dict[str, Any]]:
    data = trace.model_dump()
    run_id = state.get("run_metadata", {}).get("run_id")
    if run_id:
        append_stage_trace(str(run_id), data)
    return list(state.get("stage_traces", [])) + [data]


def disabled_update(state: dict[str, Any], agent_key: str) -> dict[str, Any]:
    agent = get_agent(agent_key)
    trace = StageTrace(
        iteration=int(state.get("iterations", 1) or 1),
        stage_index=stage_index(state),
        agent_key=agent_key,
        agent_name=agent.name,
        status="skipped",
        output_summary=["Agent execution mode is Disabled."],
        agent_explanation="The agent was skipped by configuration.",
    )
    return {
        "stage_traces": record_trace(state, trace),
        "working_memory": append_timeline(state, f"{agent.name} skipped."),
    }


def failure_update(
    state: dict[str, Any],
    agent_key: str,
    started_ms: int,
    error: Exception,
) -> dict[str, Any]:
    agent = get_agent(agent_key)
    error_data = {"agent_key": agent_key, "message": str(error), "timestamp_ms": now_ms()}
    trace = StageTrace(
        iteration=int(state.get("iterations", 1) or 1),
        stage_index=stage_index(state),
        agent_key=agent_key,
        agent_name=agent.name,
        status="failed",
        output_summary=[str(error)],
        decision_basis=["The node raised an exception or returned invalid structured output."],
        agent_explanation="The agent failed before producing a valid schema-conformant result.",
        duration_ms=now_ms() - started_ms,
    )
    return {
        "errors": list(state.get("errors", [])) + [error_data],
        "stage_traces": record_trace(state, trace),
        "working_memory": append_timeline(state, f"{agent.name} failed: {error}"),
    }