"""Middleware — auto-instrument KOAN agents with decision tracing."""

from __future__ import annotations

import uuid
from typing import TYPE_CHECKING

from koan.graph.tracer import DecisionTracer

if TYPE_CHECKING:
    from koan.agents.base import BaseAgent
    from koan.agents.types import AgentResult, RunContext
    from koan.graph.client import GraphClient


async def traced_agent(
    agent: BaseAgent,
    input_text: str,
    *,
    client: GraphClient,
    context: RunContext | None = None,
    thread_id: str | None = None,
) -> AgentResult:
    """Run an agent with automatic decision tracing.

    Creates a DecisionTracer for the run, records a routing decision
    and the final outcome, then returns the agent's result.

    Usage::

        result = await traced_agent(
            my_agent, "Hello!",
            client=graph_client,
            context=run_context,
        )
    """
    run_id = str(uuid.uuid4())
    tid = thread_id or (context.thread_id if context else None) or str(uuid.uuid4())

    async with DecisionTracer(
        client,
        run_id=run_id,
        agent_id=agent.name,
        thread_id=tid,
    ) as tracer:
        # Record initial routing decision
        decision_id = await tracer.record_decision(
            description=f"Agent '{agent.name}' processing input",
            decision_type="routing",
            confidence=1.0,
            reasoning="Direct agent invocation via traced_agent()",
        )

        # Run the actual agent
        result = await agent.run(input_text, context=context)

        # Record tool call actions if any were made
        if result.tool_calls_made > 0:
            await tracer.record_action(
                decision_id=decision_id,
                action_type="tool_call",
                tool_name=f"{result.tool_calls_made} tool call(s)",
            )

        # Record the outcome
        status = "failure" if result.is_error else "success"
        await tracer.record_outcome(status=status, result=result.output)

        # Attach trace metadata to result
        result.metadata["run_id"] = run_id
        result.metadata["thread_id"] = tid

    return result
