"""Multi-agent feedback propagation across handoff chains."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from uuid import uuid4

import structlog

if TYPE_CHECKING:
    from koan.graph.client import GraphClient

log = structlog.get_logger()


@dataclass
class PropagationResult:
    """Result of propagating feedback across a chain."""

    source_run_id: str
    propagated_to: list[str] = field(default_factory=list)
    feedback_ids: list[str] = field(default_factory=list)


async def link_runs(
    client: GraphClient,
    *,
    parent_run_id: str,
    child_run_id: str,
    relationship: str = "DELEGATED_TO",
) -> None:
    """Link two runs in the graph to form a handoff chain.

    Used by orchestrators, sequential chains, and supervisors to record
    that one run delegated to another.
    """
    await client.execute_write(
        f"""
        MATCH (parent:Run {{run_id: $parent_id}})
        MATCH (child:Run {{run_id: $child_id}})
        CREATE (parent)-[:{relationship}]->(child)
        """,
        {"parent_id": parent_run_id, "child_id": child_run_id},
    )


class FeedbackPropagator:
    """Propagates feedback signals across multi-agent handoff chains.

    When a user gives feedback on the final output of a chain (e.g.,
    orchestrator → agent A → agent B), this propagator walks the
    ``DELEGATED_TO`` edges backwards and creates derived feedback
    for each upstream run.

    The propagated score is attenuated by a decay factor at each hop
    to reflect that upstream agents are less directly responsible.

    Usage::

        propagator = FeedbackPropagator(graph_client)
        result = await propagator.propagate(
            run_id="final-run-id",
            score=1.0,
            comment="Great answer!",
        )
        # result.propagated_to = ["parent-run", "grandparent-run"]
    """

    def __init__(
        self,
        client: GraphClient,
        *,
        decay_factor: float = 0.7,
        max_hops: int = 5,
    ) -> None:
        self._client = client
        self._decay = decay_factor
        self._max_hops = max_hops

    async def propagate(
        self,
        *,
        run_id: str,
        score: float,
        feedback_type: str = "propagated",
        comment: str = "",
    ) -> PropagationResult:
        """Propagate feedback from a run to all upstream runs in the chain.

        Args:
            run_id: The run that received direct feedback.
            score: The feedback score (-1.0 to 1.0).
            feedback_type: Type label for the propagated feedback.
            comment: Optional comment from the original feedback.

        Returns:
            PropagationResult with list of runs that received propagated feedback.
        """
        result = PropagationResult(source_run_id=run_id)

        # Walk backwards through DELEGATED_TO edges
        upstream_runs = await self._find_upstream_runs(run_id)

        current_score = score
        for i, upstream_id in enumerate(upstream_runs):
            current_score *= self._decay
            fid = await self._record_propagated_feedback(
                run_id=upstream_id,
                score=round(current_score, 3),
                feedback_type=feedback_type,
                comment=(
                    f"Propagated (hop {i + 1}): {comment}"
                    if comment
                    else f"Propagated from {run_id[:8]} (hop {i + 1})"
                ),
                source_run_id=run_id,
            )
            result.propagated_to.append(upstream_id)
            result.feedback_ids.append(fid)

            log.debug(
                "feedback.propagated",
                source=run_id[:8],
                target=upstream_id[:8],
                hop=i + 1,
                score=current_score,
            )

        return result

    async def _find_upstream_runs(self, run_id: str) -> list[str]:
        """Find all runs that delegated (directly or transitively) to this run."""
        records = await self._client.execute(
            """
            MATCH path = (ancestor:Run)-[:DELEGATED_TO*1..]->(target:Run {run_id: $run_id})
            UNWIND nodes(path) AS n
            WITH DISTINCT n
            WHERE n.run_id <> $run_id
            RETURN n.run_id AS run_id
            """,
            {"run_id": run_id},
        )
        # Limit hops
        return [r.data["run_id"] for r in records[:self._max_hops]]

    async def _record_propagated_feedback(
        self,
        *,
        run_id: str,
        score: float,
        feedback_type: str,
        comment: str,
        source_run_id: str,
    ) -> str:
        """Record a propagated feedback node."""
        from datetime import UTC, datetime  # noqa: PLC0415

        feedback_id = str(uuid4())
        await self._client.execute_write(
            """
            MATCH (r:Run {run_id: $run_id})
            CREATE (f:Feedback {
                feedback_id: $feedback_id,
                feedback_type: $feedback_type,
                score: $score,
                comment: $comment,
                timestamp: $timestamp,
                source_run_id: $source_run_id,
                propagated: true
            })
            CREATE (r)-[:HAS_FEEDBACK]->(f)
            """,
            {
                "run_id": run_id,
                "feedback_id": feedback_id,
                "feedback_type": feedback_type,
                "score": score,
                "comment": comment,
                "timestamp": datetime.now(UTC).isoformat(),
                "source_run_id": source_run_id,
            },
        )
        return feedback_id

    async def get_chain(self, run_id: str) -> list[dict[str, Any]]:
        """Get the full delegation chain for a run (for debugging/visualization).

        Returns list of dicts with run_id, agent_id, status for each run in the chain.
        """
        records = await self._client.execute(
            """
            MATCH path = (root:Run)-[:DELEGATED_TO*0..]->(target:Run {run_id: $run_id})
            UNWIND nodes(path) AS n
            MATCH (a:Agent)-[:EXECUTED]->(n)
            RETURN DISTINCT n.run_id AS run_id,
                            a.agent_id AS agent_id,
                            n.status AS status
            """,
            {"run_id": run_id},
        )
        return [r.data for r in records]
