"""GraphContextProvider — query the decision graph to build agent context."""

from __future__ import annotations

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

if TYPE_CHECKING:
    from koan.graph.client import GraphClient
    from koan.graph.feedback import ConfidenceCalibrator


@dataclass
class GraphContext:
    """Context gathered from the decision graph for agent augmentation."""

    entity_history: list[dict[str, Any]] = field(default_factory=list)
    precedents: list[dict[str, Any]] = field(default_factory=list)
    domain_profile: dict[str, Any] = field(default_factory=dict)
    active_policies: list[dict[str, Any]] = field(default_factory=list)
    feedback_summary: dict[str, Any] = field(default_factory=dict)

    def to_prompt(self) -> str:
        """Format the graph context as a text block for injection into a system prompt."""
        sections: list[str] = []

        if self.domain_profile:
            total = self.domain_profile.get("total_decisions", 0)
            types = self.domain_profile.get("decision_types", {})
            outcomes = self.domain_profile.get("outcomes", {})
            sections.append(
                f"[Domain Profile] You have made {total} past decisions. "
                f"Decision types: {types}. Outcome summary: {outcomes}."
            )

        if self.active_policies:
            policy_lines: list[str] = []
            for p in self.active_policies:
                name = p.get("name", p.get("policy_id", "?"))
                desc = p.get("description", "")
                rules = p.get("rules", [])
                line = f"  - {name}: {desc}"
                if rules:
                    line += "\n    Rules: " + "; ".join(str(r) for r in rules)
                policy_lines.append(line)
            sections.append(
                "[Active Policies — you MUST follow these rules]\n"
                + "\n".join(policy_lines)
            )

        if self.entity_history:
            entity_lines: list[str] = []
            for eh in self.entity_history[:5]:
                eid = eh.get("entity_id", "?")
                etype = eh.get("entity_type", "?")
                count = eh.get("decision_count", 0)
                last_outcome = eh.get("last_outcome", "unknown")
                entity_lines.append(
                    f"  - {etype} '{eid}': {count} prior decision(s), last outcome: {last_outcome}"
                )
            sections.append("[Entity History]\n" + "\n".join(entity_lines))

        if self.precedents:
            prec_lines: list[str] = []
            for p in self.precedents[:5]:
                desc = p.get("description", "?")
                outcome = p.get("outcome_status", "unknown")
                reasoning = p.get("reasoning", "")
                tool = p.get("tool_name")
                feedback_score = p.get("feedback_score")
                feedback_type = p.get("feedback_type")

                tool_str = f", tool: {tool}" if tool else ""

                # Build quality indicator from outcome + feedback
                quality = f"[{outcome}]"
                if feedback_score is not None and feedback_score != 0:
                    emoji = "+" if feedback_score > 0 else "-"
                    fb_label = feedback_type or ("liked" if feedback_score > 0 else "disliked")
                    quality = f"[{outcome}, {fb_label} ({emoji}{feedback_score})]"

                line = f"  - {quality} {desc}{tool_str}"
                if reasoning:
                    line += f"\n    Reasoning: {reasoning}"
                prec_lines.append(line)
            sections.append(
                "[Relevant Precedents — learn from successful decisions, avoid patterns from failed ones]\n"
                + "\n".join(prec_lines)
            )

        if self.feedback_summary:
            total = self.feedback_summary.get("total_feedback", 0)
            sat_rate = self.feedback_summary.get("satisfaction_rate", 0.0)
            avg = self.feedback_summary.get("avg_score", 0.0)
            comments = self.feedback_summary.get("recent_comments", [])
            fb_lines = [
                f"  Past feedback: {total} ratings, {sat_rate:.0%} satisfaction, avg score: {avg}",
            ]
            if comments:
                fb_lines.append("  Recent comments: " + "; ".join(comments[:3]))
            sections.append("[Feedback History]\n" + "\n".join(fb_lines))

        if not sections:
            return ""

        return (
            "\n\n--- GRAPH CONTEXT (from past agent decisions) ---\n"
            + "\n\n".join(sections)
            + "\n--- END GRAPH CONTEXT ---\n"
        )


def _precedent_quality_key(prec: dict[str, Any]) -> tuple[int, float]:
    """Sort key for precedents: successful+liked first, failed/disliked last.

    Returns (outcome_rank, feedback_score) so higher is better.
    """
    outcome = prec.get("outcome_status", "unknown")
    outcome_rank = {"success": 2, "unknown": 1, "failure": 0}.get(outcome, 1)
    feedback = prec.get("feedback_score") or 0.0
    return (outcome_rank, float(feedback))


class GraphContextProvider:
    """Queries the decision graph to build context for agent augmentation.

    Before an agent run, call ``gather()`` with entity IDs and/or the agent name
    to retrieve relevant precedents, entity history, domain profile, policies,
    and feedback history.

    Usage::

        provider = GraphContextProvider(graph_client)
        ctx = await provider.gather(
            agent_id="finance-agent",
            entity_ids=["acct-001"],
        )
        prompt_block = ctx.to_prompt()
    """

    def __init__(
        self,
        client: GraphClient,
        *,
        calibrator: ConfidenceCalibrator | None = None,
    ) -> None:
        self._client = client
        self._calibrator = calibrator

    async def gather(
        self,
        *,
        agent_id: str | None = None,
        entity_ids: list[str] | None = None,
        decision_type: str | None = None,
        limit: int = 5,
    ) -> GraphContext:
        """Gather context from the graph.

        Args:
            agent_id: Fetch domain profile for this agent.
            entity_ids: Fetch history and precedents for these entities.
            decision_type: Filter precedents to this type.
            limit: Max precedents/history entries.

        Returns:
            GraphContext with populated fields.
        """
        ctx = GraphContext()

        if agent_id:
            ctx.domain_profile = await self._get_domain_profile(agent_id)

        if entity_ids:
            ctx.entity_history = await self._get_entity_history(entity_ids, limit=limit)

            if decision_type:
                ctx.precedents = await self._get_precedents(
                    entity_ids, decision_type=decision_type, limit=limit,
                )

        # If no entity-specific precedents found, fall back to cross-domain
        if not ctx.precedents and decision_type:
            ctx.precedents = await self._get_cross_domain_precedents(
                decision_type=decision_type, limit=limit,
            )

        ctx.active_policies = await self._get_active_policies(limit=limit)

        if self._calibrator:
            ctx.feedback_summary = await self._calibrator.get_feedback_summary(
                agent_id=agent_id,
                decision_type=decision_type,
            )

        return ctx

    async def _get_domain_profile(self, agent_id: str) -> dict[str, Any]:
        """Get aggregate stats for an agent's past decisions."""
        records = await self._client.execute(
            """
            MATCH (a:Agent {agent_id: $agent_id})-[:EXECUTED]->(r:Run)-[:CONTAINS]->(d:Decision)
            RETURN count(d) AS total_decisions
            """,
            {"agent_id": agent_id},
        )

        if not records:
            return {}

        row = records[0].data
        total = row.get("total_decisions", 0)

        # Get outcome distribution
        outcome_records = await self._client.execute(
            """
            MATCH (a:Agent {agent_id: $agent_id})-[:EXECUTED]->(r:Run)-[:PRODUCED]->(o:Outcome)
            RETURN o.status AS status, count(o) AS count
            """,
            {"agent_id": agent_id},
        )

        outcomes: dict[str, int] = {}
        for r in outcome_records:
            status = r.data.get("status", "unknown")
            count = r.data.get("count", 0)
            outcomes[status] = count

        # Build type distribution
        type_records = await self._client.execute(
            """
            MATCH (a:Agent {agent_id: $agent_id})-[:EXECUTED]->(r:Run)-[:CONTAINS]->(d:Decision)
            RETURN d.decision_type AS dtype, count(d) AS count
            """,
            {"agent_id": agent_id},
        )
        type_dist: dict[str, int] = {}
        for r in type_records:
            dtype = r.data.get("dtype", "unknown")
            count = r.data.get("count", 0)
            type_dist[dtype] = count

        return {
            "total_decisions": total,
            "decision_types": type_dist,
            "outcomes": outcomes,
        }

    async def _get_entity_history(
        self, entity_ids: list[str], *, limit: int = 5,
    ) -> list[dict[str, Any]]:
        """Get decision history for each entity."""
        results: list[dict[str, Any]] = []

        for eid in entity_ids:
            records = await self._client.execute(
                """
                MATCH (d:Decision)-[:REFERENCED]->(e:Entity {entity_id: $entity_id})
                OPTIONAL MATCH (d)<-[:CONTAINS]-(r:Run)-[:PRODUCED]->(o:Outcome)
                WITH e.entity_id AS entity_id,
                     e.entity_type AS entity_type,
                     count(d) AS decision_count,
                     collect(o.status)[-1] AS last_outcome
                RETURN entity_id, entity_type, decision_count, last_outcome
                LIMIT 1
                """,
                {"entity_id": eid},
            )

            if records:
                results.append(records[0].data)

        return results[:limit]

    async def _get_precedents(
        self,
        entity_ids: list[str],
        *,
        decision_type: str,
        limit: int = 5,
    ) -> list[dict[str, Any]]:
        """Get precedent decisions for entities of a given type.

        Joins to Feedback nodes so each precedent carries a feedback score,
        letting the LLM distinguish good decisions from bad ones.
        """
        all_precs: list[dict[str, Any]] = []

        for eid in entity_ids:
            records = await self._client.execute(
                """
                MATCH (d:Decision)-[:REFERENCED]->(e:Entity {entity_id: $entity_id})
                WHERE d.decision_type = $decision_type
                OPTIONAL MATCH (d)<-[:CONTAINS]-(r:Run)-[:PRODUCED]->(o:Outcome)
                OPTIONAL MATCH (d)-[:LED_TO]->(a:Action)
                OPTIONAL MATCH (r)-[:HAS_FEEDBACK]->(f:Feedback)
                RETURN d.decision_id AS decision_id,
                       d.description AS description,
                       d.reasoning AS reasoning,
                       d.confidence AS confidence,
                       d.timestamp AS timestamp,
                       o.status AS outcome_status,
                       a.tool_name AS tool_name,
                       f.score AS feedback_score,
                       f.feedback_type AS feedback_type
                ORDER BY d.timestamp DESC
                LIMIT $limit
                """,
                {"entity_id": eid, "decision_type": decision_type, "limit": limit},
            )

            all_precs.extend(r.data for r in records)

        # Sort: successful+liked first, then neutral, then failed/disliked
        return sorted(all_precs[:limit], key=_precedent_quality_key, reverse=True)

    async def _get_cross_domain_precedents(
        self,
        *,
        decision_type: str,
        limit: int = 5,
    ) -> list[dict[str, Any]]:
        """Get precedents by decision_type across all entities.

        Falls back to this when no entity-specific precedents exist,
        so the agent still benefits from past experience with the same
        type of decision even for new/unknown entities.
        """
        records = await self._client.execute(
            """
            MATCH (d:Decision)
            WHERE d.decision_type = $decision_type
            OPTIONAL MATCH (d)<-[:CONTAINS]-(r:Run)-[:PRODUCED]->(o:Outcome)
            OPTIONAL MATCH (d)-[:LED_TO]->(a:Action)
            OPTIONAL MATCH (d)-[:REFERENCED]->(e:Entity)
            OPTIONAL MATCH (r)-[:HAS_FEEDBACK]->(f:Feedback)
            RETURN d.decision_id AS decision_id,
                   d.description AS description,
                   d.reasoning AS reasoning,
                   d.confidence AS confidence,
                   d.timestamp AS timestamp,
                   o.status AS outcome_status,
                   a.tool_name AS tool_name,
                   collect(DISTINCT e.entity_id) AS entity_ids,
                   f.score AS feedback_score,
                   f.feedback_type AS feedback_type
            ORDER BY d.timestamp DESC
            LIMIT $limit
            """,
            {"decision_type": decision_type, "limit": limit},
        )
        precs = [r.data for r in records]
        return sorted(precs, key=_precedent_quality_key, reverse=True)

    async def _get_active_policies(self, *, limit: int = 10) -> list[dict[str, Any]]:
        """Get all policies defined in the graph, including rules."""
        records = await self._client.execute(
            """
            MATCH (p:Policy)
            RETURN p.policy_id AS policy_id,
                   p.name AS name,
                   p.description AS description,
                   p.version AS version,
                   p.source AS source,
                   p.rules AS rules
            LIMIT $limit
            """,
            {"limit": limit},
        )
        return [r.data for r in records]
