"""GraphAugmentedAgent — wraps an agent with graph-informed reasoning."""

from __future__ import annotations

import uuid
from typing import TYPE_CHECKING

import structlog

from koan.agents.base import BaseAgent
from koan.agents.types import AgentConfig, AgentResult, RunContext
from koan.graph.context import GraphContext, GraphContextProvider
from koan.graph.feedback import ConfidenceCalibrator
from koan.graph.gated_tools import PolicyGatedTools
from koan.graph.policies import PolicyManager
from koan.graph.tracer import DecisionTracer

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.agents.events import AgentEvent
    from koan.graph.client import GraphClient
    from koan.llm.client import LLMClient
    from koan.tools.types import ToolProvider

log = structlog.get_logger()


class GraphAugmentedAgent(BaseAgent):
    """An agent that reads from the decision graph before reasoning and writes back after.

    Closes the read-write loop:
    1. **Before run**: Queries the graph for domain profile, entity history,
       precedents, and active policies via ``GraphContextProvider``.
    2. **During run**: Injects the graph context into the system prompt so
       the LLM reasons with awareness of past decisions.
    3. **After run**: Records the decision trace via ``DecisionTracer``.

    Usage::

        agent = GraphAugmentedAgent(
            "finance-bot",
            client=llm_client,
            config=AgentConfig(provider="openai", model="gpt-4o"),
            tools=registry,
            graph_client=graph,
            entity_ids=["acct-001"],
        )
        result = await agent.run("Process the billing inquiry for acct-001")
        # Agent saw past decisions for acct-001 and wrote new trace back
    """

    def __init__(
        self,
        name: str,
        *,
        client: LLMClient,
        config: AgentConfig | None = None,
        tools: ToolProvider | None = None,
        graph_client: GraphClient,
        inner_agent: BaseAgent | None = None,
        entity_ids: list[str] | None = None,
        decision_type: str | None = None,
        context_limit: int = 5,
    ) -> None:
        super().__init__(name, client=client, config=config, tools=tools)
        self._graph = graph_client
        self._calibrator = ConfidenceCalibrator(graph_client)
        self._context_provider = GraphContextProvider(
            graph_client, calibrator=self._calibrator,
        )
        self._entity_ids = entity_ids or []
        self._decision_type = decision_type
        self._context_limit = context_limit
        self._inner = inner_agent
        self._last_run_id: str | None = None
        self._policy_manager = PolicyManager(graph_client)

    async def _gather_context(self) -> GraphContext:
        """Query the graph for relevant context."""
        return await self._context_provider.gather(
            agent_id=self.name,
            entity_ids=self._entity_ids if self._entity_ids else None,
            decision_type=self._decision_type,
            limit=self._context_limit,
        )

    def _augment_prompt(self, graph_context: GraphContext) -> str:
        """Combine the agent's system prompt with graph context."""
        base_prompt = self.config.system_prompt
        context_block = graph_context.to_prompt()

        if not context_block:
            return base_prompt

        if base_prompt:
            return base_prompt + "\n" + context_block
        return context_block.strip()

    def _gate_tools(self, graph_context: GraphContext) -> ToolProvider | None:
        """Wrap tools with policy-based gating if policies constrain tool usage.

        Returns a PolicyGatedTools wrapper if policies specify tool constraints,
        or the original tools if no constraints exist (permissive cold-start).
        """
        if not self.tools or not graph_context.active_policies:
            return self.tools

        allowed = PolicyGatedTools.extract_tools_from_policies(
            graph_context.active_policies,
        )

        if allowed is None:
            return self.tools  # No tool constraints in policies

        log.info(
            "graph_augmented.tools_gated",
            allowed=allowed,
            total_tools=len(self.tools.definitions()) if self.tools else 0,
        )
        return PolicyGatedTools(self.tools, allowed_tools=allowed)

    def _get_inner_agent(
        self,
        augmented_prompt: str,
        *,
        tools: ToolProvider | None = None,
    ) -> BaseAgent:
        """Get or create the inner agent with the augmented prompt."""
        effective_tools = tools or self.tools

        if self._inner is not None:
            # Update inner agent's system prompt and tools
            self._inner.config = AgentConfig(
                provider=self._inner.config.provider,
                model=self._inner.config.model,
                temperature=self._inner.config.temperature,
                max_tokens=self._inner.config.max_tokens,
                max_iterations=self._inner.config.max_iterations,
                system_prompt=augmented_prompt,
            )
            if effective_tools:
                self._inner.tools = effective_tools  # type: ignore[assignment]
            return self._inner

        # Create a ReactAgent on the fly
        from koan.agents.react import ReactAgent  # noqa: PLC0415

        return ReactAgent(
            self.name,
            client=self.client,
            config=AgentConfig(
                provider=self.config.provider,
                model=self.config.model,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                max_iterations=self.config.max_iterations,
                system_prompt=augmented_prompt,
            ),
            tools=effective_tools,  # type: ignore[arg-type]
        )

    async def run(self, input_text: str, *, context: RunContext | None = None) -> AgentResult:
        """Execute with graph-augmented reasoning.

        1. Gather context from the graph
        2. Augment the system prompt
        3. Run the inner agent
        4. Record the decision trace
        """
        run_id = str(uuid.uuid4())
        tid = (context.thread_id if context else None) or str(uuid.uuid4())

        # Step 1: Read from graph
        graph_context = await self._gather_context()

        log.debug(
            "graph_augmented.context_gathered",
            agent=self.name,
            entities=len(graph_context.entity_history),
            precedents=len(graph_context.precedents),
            policies=len(graph_context.active_policies),
            has_profile=bool(graph_context.domain_profile),
        )

        # Step 2: Augment prompt and gate tools
        augmented_prompt = self._augment_prompt(graph_context)
        gated_tools = self._gate_tools(graph_context)

        # Step 3: Run inner agent with policy-gated tools
        inner = self._get_inner_agent(augmented_prompt, tools=gated_tools)

        # Step 4: Trace the run
        async with DecisionTracer(
            self._graph,
            run_id=run_id,
            agent_id=self.name,
            thread_id=tid,
        ) as tracer:
            result = await inner.run(input_text, context=context)

            # Build a meaningful description from the agent's actual output
            # so precedents show what the agent *did*, not what the user asked
            output_summary = result.output[:200] if result.output else "No output"
            dtype = self._decision_type or "augmented"

            decision_id = await tracer.record_decision(
                description=f"Q: {input_text[:80]} → A: {output_summary}",
                decision_type=dtype,
                confidence=result.metadata.get("confidence", 1.0),
                reasoning=(
                    f"Used {result.tool_calls_made} tool(s). "
                    f"Context: {len(graph_context.precedents)} precedents, "
                    f"{len(graph_context.active_policies)} policies, "
                    f"{len(graph_context.entity_history)} entities."
                ),
            )

            # Record tool actions
            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)",
                )

            # Link entities
            for eid in self._entity_ids:
                await tracer.link_entity(
                    decision_id=decision_id,
                    entity_id=eid,
                    entity_type="referenced",
                )

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

            # Link decision to matching policies (creates GOVERNED_BY edges)
            linked_policies = await self._policy_manager.link_decision_to_matching_policies(
                decision_id=decision_id, decision_type=dtype,
            )

        self._last_run_id = run_id

        # Evolve: promote new patterns, retire underperforming policies
        promoted = await self._policy_manager.auto_evolve()
        retired = await self._policy_manager.retire_underperforming()

        result.metadata["run_id"] = run_id
        result.metadata["thread_id"] = tid
        result.metadata["graph_context"] = {
            "precedents": len(graph_context.precedents),
            "entities": len(graph_context.entity_history),
            "policies": len(graph_context.active_policies),
            "has_feedback": bool(graph_context.feedback_summary),
            "linked_policies": linked_policies,
            "new_policies": promoted,
            "retired_policies": retired,
        }

        return result

    async def run_stream(
        self, input_text: str, *, context: RunContext | None = None,
    ) -> AsyncIterator[AgentEvent]:
        """Stream with graph-augmented reasoning.

        Gathers graph context, augments the prompt, then delegates
        to the inner agent's ``run_stream()``.
        """
        from koan.agents.events import (  # noqa: PLC0415
            AgentComplete,
            AgentError,
            ToolCallEnd,
            ToolCallStart,
        )
        from koan.agents.stream import StreamingReactAgent  # noqa: PLC0415

        run_id = str(uuid.uuid4())
        tid = (context.thread_id if context else None) or str(uuid.uuid4())

        # Gather, augment, and gate tools
        graph_context = await self._gather_context()
        augmented_prompt = self._augment_prompt(graph_context)
        gated_tools = self._gate_tools(graph_context)

        # Create a streaming inner agent with policy-gated tools
        streaming_agent = StreamingReactAgent(
            self.name,
            client=self.client,
            config=AgentConfig(
                provider=self.config.provider,
                model=self.config.model,
                temperature=self.config.temperature,
                max_tokens=self.config.max_tokens,
                max_iterations=self.config.max_iterations,
                system_prompt=augmented_prompt,
            ),
            tools=gated_tools,  # type: ignore[arg-type]
        )

        # Stream events, capture final result and individual tool calls for tracing
        final_output = ""
        tool_calls_count = 0
        tool_call_details: list[dict[str, str]] = []

        gc_meta = {
            "precedents": len(graph_context.precedents),
            "entities": len(graph_context.entity_history),
            "policies": len(graph_context.active_policies),
            "has_feedback": bool(graph_context.feedback_summary),
        }

        async for event in streaming_agent.run_stream(input_text, context=context):
            if isinstance(event, ToolCallStart):
                tool_call_details.append({"tool_name": event.tool_name, "result": ""})
            elif isinstance(event, ToolCallEnd):
                # Update the last tool call with its result
                if tool_call_details:
                    tool_call_details[-1]["result"] = event.result or ""
            elif isinstance(event, AgentComplete):
                final_output = event.output
                tool_calls_count = event.tool_calls_made
                event.metadata["run_id"] = run_id
                event.metadata["graph_context"] = gc_meta
            elif isinstance(event, AgentError):
                final_output = event.output or ""
            yield event

        self._last_run_id = run_id

        # Write trace after streaming completes
        dtype = self._decision_type or "augmented"
        output_summary = final_output[:200] if final_output else "No output"
        tools_used = [tc["tool_name"] for tc in tool_call_details] if tool_call_details else []

        async with DecisionTracer(
            self._graph, run_id=run_id, agent_id=self.name, thread_id=tid,
        ) as tracer:
            decision_id = await tracer.record_decision(
                description=f"Q: {input_text[:80]} → A: {output_summary}",
                decision_type=dtype,
                confidence=1.0,
                reasoning=(
                    f"Tools: {', '.join(tools_used) or 'none'}. "
                    f"Context: {len(graph_context.precedents)} precedents, "
                    f"{len(graph_context.active_policies)} policies."
                ),
            )

            # Record each tool call individually for granular tracing
            if tool_call_details:
                for tc in tool_call_details:
                    await tracer.record_action(
                        decision_id=decision_id,
                        action_type="tool_call",
                        tool_name=tc["tool_name"],
                    )
            elif tool_calls_count > 0:
                # Fallback: summary if details weren't captured
                await tracer.record_action(
                    decision_id=decision_id,
                    action_type="tool_call",
                    tool_name=f"{tool_calls_count} tool call(s)",
                )

            for eid in self._entity_ids:
                await tracer.link_entity(
                    decision_id=decision_id, entity_id=eid, entity_type="referenced",
                )

            status = "success" if final_output else "failure"
            await tracer.record_outcome(status=status, result=final_output[:500])

            # Link decision to matching policies
            await self._policy_manager.link_decision_to_matching_policies(
                decision_id=decision_id, decision_type=dtype,
            )

        # Evolve: promote new patterns, retire underperforming policies
        await self._policy_manager.auto_evolve()
        await self._policy_manager.retire_underperforming()

    async def submit_feedback(
        self,
        *,
        run_id: str | None = None,
        satisfied: bool,
        comment: str = "",
    ) -> str:
        """Submit user satisfaction feedback for a run.

        Args:
            run_id: The run to attach feedback to. Defaults to the last run.
            satisfied: Whether the user was satisfied with the answer.
            comment: Optional free-text comment.

        Returns:
            The feedback_id.
        """
        from koan.graph.feedback import FeedbackType, record_feedback  # noqa: PLC0415

        rid = run_id or self._last_run_id
        if not rid:
            msg = "No run_id provided and no previous run to attach feedback to."
            raise ValueError(msg)

        score = 1.0 if satisfied else -1.0
        ftype = FeedbackType.THUMBS_UP if satisfied else FeedbackType.THUMBS_DOWN

        return await record_feedback(
            self._graph,
            run_id=rid,
            feedback_type=ftype,
            score=score,
            comment=comment,
        )
