context_agent.py

from __future__ import annotations

import asyncio
import logging
from typing import Any, Awaitable, Callable, Optional

from backend.agents.base import BaseAgent
from backend.models.enums import AgentPhase, ContextSource
from backend.models.schemas import (
    ClarifiedIntent,
    ContextItem,
    ExecutionPlan,
    ResourcePlan,
    SessionState,
    StepContext,
    ToolRequirement,
)
from backend.prompts.context_prompt import (
    CONTEXT_ANALYSIS_PROMPT,
    CONTEXT_SUMMARY_PROMPT,
)
from backend.services.exa_service import ExaService
from backend.services.llm_service import LLMService
from backend.services.pageindex_service import PageIndexService

logger = logging.getLogger(__name__)


class ContextGatheringAgent(BaseAgent):
    """Phase 4.5 agent: multi-source context gathering and resource planning.

    Three-step process:
      1. Analyze all plan steps and determine what retrieval each step needs.
      2. Execute retrievals in parallel: PageIndex (KB), Exa (web), prior phases.
      3. Summarize and assemble a structured ResourcePlan.
    """

    phase = AgentPhase.CONTEXT_GATHERING
    name = "Context Gathering Agent"

    def __init__(
        self,
        llm_service: LLMService,
        exa_service: ExaService,
        pageindex_service: PageIndexService,
    ) -> None:
        super().__init__(llm_service)
        self.exa = exa_service
        self.pageindex = pageindex_service

    def get_system_prompt(self) -> str:
        return CONTEXT_ANALYSIS_PROMPT

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def process(
        self,
        input_data: Any,
        on_stream: Optional[Callable[[str], Awaitable[None]]] = None,
    ) -> ResourcePlan:
        """Convenience entry — expects a dict with plan, clarified_intent, session_state."""
        plan: ExecutionPlan = input_data["plan"]
        clarified_intent: ClarifiedIntent = input_data["clarified_intent"]
        session_state: SessionState = input_data["session_state"]
        on_step_update: Optional[Callable] = input_data.get("on_step_update")
        return await self.gather(plan, clarified_intent, session_state, on_step_update)

    async def gather(
        self,
        plan: ExecutionPlan,
        clarified_intent: ClarifiedIntent,
        session_state: SessionState,
        on_step_update: Optional[Callable[[int, str, dict], Awaitable[None]]] = None,
    ) -> ResourcePlan:
        """Full context-gathering pipeline.

        Args:
            plan:             The execution plan from Phase 4.
            clarified_intent: The clarified intent from Phase 3.
            session_state:    Full session state for prior-phase extraction.
            on_step_update:   Optional async callback(step_number, status, data).
        """
        kb_docs = list(self.pageindex.indexed_documents.keys()) if self.pageindex.available else []

        # --- Step 1: Analyse steps and derive retrieval requirements ---
        requirements_per_step = await self._analyse_steps(plan, clarified_intent, kb_docs)

        # --- Step 2: Execute retrievals in parallel per step ---
        step_contexts: list[StepContext] = []
        retrieval_tasks = []

        for req in requirements_per_step:
            step_num = req["step_number"]
            step_desc = next(
                (s.description for s in plan.steps if s.step_number == step_num), ""
            )
            tool_reqs = [ToolRequirement(**t) for t in req.get("tools", [])]
            retrieval_tasks.append(
                self._gather_step_context(
                    step_num, step_desc, tool_reqs, session_state, on_step_update
                )
            )

        step_contexts = list(await asyncio.gather(*retrieval_tasks, return_exceptions=False))
        # Filter out any exceptions (graceful degradation)
        step_contexts = [sc for sc in step_contexts if isinstance(sc, StepContext)]

        # --- Step 3: Summarize and build ResourcePlan ---
        resource_plan = await self._build_resource_plan(plan, step_contexts)

        logger.info(
            "ContextGatheringAgent: gathered contexts for %d steps "
            "(KB docs=%d, web=%d)",
            len(step_contexts),
            len(resource_plan.knowledge_base_docs_used),
            resource_plan.web_searches_performed,
        )
        return resource_plan

    # ------------------------------------------------------------------
    # Step 1: Analyse retrieval needs
    # ------------------------------------------------------------------

    async def _analyse_steps(
        self,
        plan: ExecutionPlan,
        clarified_intent: ClarifiedIntent,
        kb_docs: list[str],
    ) -> list[dict]:
        """Single LLM call to determine retrieval requirements for all steps."""
        kb_section = (
            f"Available knowledge base documents: {kb_docs}"
            if kb_docs
            else "Knowledge base: no documents currently indexed"
        )
        user_msg = (
            f"Execution plan:\n{plan.model_dump_json(indent=2)}\n\n"
            f"Clarified intent:\n{clarified_intent.model_dump_json(indent=2)}\n\n"
            f"{kb_section}\n\n"
            "Analyse each step and produce the retrieval requirements JSON."
        )
        try:
            result = await self.llm.generate_json(CONTEXT_ANALYSIS_PROMPT, user_msg)
            return result.get("step_requirements", [])
        except Exception:
            logger.exception("ContextGatheringAgent: step analysis LLM call failed")
            # Return a fallback: prior_context only for each step
            return [
                {
                    "step_number": s.step_number,
                    "tools": [
                        {
                            "tool_name": "prior_context",
                            "purpose": "Use clarified intent as context",
                            "queries": [],
                            "doc_filter": [],
                            "priority": "required",
                        }
                    ],
                }
                for s in plan.steps
            ]

    # ------------------------------------------------------------------
    # Step 2: Execute retrievals per step
    # ------------------------------------------------------------------

    async def _gather_step_context(
        self,
        step_number: int,
        step_description: str,
        tool_reqs: list[ToolRequirement],
        session_state: SessionState,
        on_step_update: Optional[Callable[[int, str, dict], Awaitable[None]]],
    ) -> StepContext:
        """Gather all context for a single step, routing to the right sources."""
        if on_step_update:
            await on_step_update(step_number, "started", {"description": step_description})

        retrieval_coroutines = []
        source_metadata: list[tuple[ToolRequirement, str]] = []

        for tool_req in tool_reqs:
            if tool_req.tool_name == "pageindex" and self.pageindex.available:
                for query in tool_req.queries:
                    retrieval_coroutines.append(
                        self._retrieve_from_kb(query, tool_req.doc_filter)
                    )
                    source_metadata.append((tool_req, query))

            elif tool_req.tool_name == "exa_search" and self.exa.available:
                queries = tool_req.queries[:3]  # cap at 3 queries per step
                if queries:
                    retrieval_coroutines.append(self._retrieve_from_exa(queries))
                    source_metadata.append((tool_req, "|".join(queries)))

            elif tool_req.tool_name == "prior_context":
                retrieval_coroutines.append(
                    self._retrieve_from_prior_phases(session_state)
                )
                source_metadata.append((tool_req, "session_state"))

            elif tool_req.tool_name == "step_dep":
                # Will be resolved at execution time — record placeholder
                dep_steps = [q for q in tool_req.queries if q.isdigit()]

                async def _dep_placeholder(d: list[str] = dep_steps) -> list[ContextItem]:
                    return [
                        ContextItem(
                            source=ContextSource.STEP_DEPENDENCY,
                            query=",".join(d),
                            content="",
                            relevance="Output will be injected at execution time",
                            source_detail=f"Depends on step(s): {', '.join(d)}",
                        )
                    ]

                retrieval_coroutines.append(_dep_placeholder())
                source_metadata.append((tool_req, f"dep:{dep_steps}"))

        raw_results: list[list[ContextItem]] = []
        if retrieval_coroutines:
            gathered = await asyncio.gather(*retrieval_coroutines, return_exceptions=True)
            for res in gathered:
                if isinstance(res, Exception):
                    logger.warning("Retrieval error in step %d: %s", step_number, res)
                elif isinstance(res, list):
                    raw_results.extend(res)

        flat_contexts: list[ContextItem] = [
            item for sublist in raw_results for item in (sublist if isinstance(sublist, list) else [sublist])
        ]

        step_ctx = StepContext(
            step_number=step_number,
            step_description=step_description,
            required_tools=tool_reqs,
            gathered_contexts=flat_contexts,
        )

        if on_step_update:
            await on_step_update(
                step_number,
                "completed",
                {
                    "description": step_description,
                    "context_count": len(flat_contexts),
                    "sources": [c.source.value for c in flat_contexts],
                },
            )

        return step_ctx

    # ------------------------------------------------------------------
    # Retrieval helpers
    # ------------------------------------------------------------------

    async def _retrieve_from_kb(
        self, query: str, doc_filter: list[str]
    ) -> list[ContextItem]:
        """Query the knowledge base via PageIndex."""
        results = await self.pageindex.search_knowledge_base(
            query, doc_filter=doc_filter or None
        )
        items: list[ContextItem] = []
        for r in results:
            if r.get("content"):
                items.append(
                    ContextItem(
                        source=ContextSource.KNOWLEDGE_BASE,
                        query=query,
                        content=r["content"],
                        relevance=f"Knowledge base: {r.get('source_doc', '')}",
                        source_detail=r.get("source_doc", ""),
                        reasoning_trace=r.get("reasoning_trace", ""),
                        metadata={
                            "doc_id": r.get("doc_id", ""),
                            "section": r.get("section", ""),
                            "confidence": r.get("confidence", ""),
                        },
                    )
                )
        return items

    async def _retrieve_from_exa(self, queries: list[str]) -> list[ContextItem]:
        """Search the web via Exa."""
        results = await self.exa.search_many(queries)
        items: list[ContextItem] = []
        for r in results:
            if r.get("snippet"):
                items.append(
                    ContextItem(
                        source=ContextSource.WEB_SEARCH,
                        query=queries[0] if queries else "",
                        content=r["snippet"],
                        relevance=r.get("title", ""),
                        source_detail=r.get("url", ""),
                        metadata={"title": r.get("title", ""), "url": r.get("url", "")},
                    )
                )
        return items

    async def _retrieve_from_prior_phases(
        self, session_state: SessionState
    ) -> list[ContextItem]:
        """Extract context from earlier pipeline phases in the session."""
        items: list[ContextItem] = []

        if session_state.user_input:
            items.append(
                ContextItem(
                    source=ContextSource.USER_INPUT,
                    query="",
                    content=session_state.user_input,
                    relevance="Original user request",
                    source_detail="user_input",
                )
            )

        if session_state.clarified_intent:
            ci = session_state.clarified_intent
            content_parts = []
            if ci.refined_goal:
                content_parts.append(f"Refined goal: {ci.refined_goal}")
            if ci.refined_constraints:
                content_parts.append(
                    "Constraints: " + "; ".join(ci.refined_constraints)
                )
            if ci.clarifications:
                pairs = "; ".join(f"{k}: {v}" for k, v in ci.clarifications.items())
                content_parts.append(f"Clarifications: {pairs}")
            if content_parts:
                items.append(
                    ContextItem(
                        source=ContextSource.PRIOR_PHASE,
                        query="",
                        content="\n".join(content_parts),
                        relevance="Clarified intent from Phase 3",
                        source_detail="clarified_intent",
                    )
                )

        if session_state.ambiguity_report:
            rep = session_state.ambiguity_report
            if rep.reasoning:
                items.append(
                    ContextItem(
                        source=ContextSource.PRIOR_PHASE,
                        query="",
                        content=rep.reasoning,
                        relevance="Ambiguity analysis from Phase 2",
                        source_detail="ambiguity_report",
                    )
                )

        return items

    # ------------------------------------------------------------------
    # Step 3: Build ResourcePlan
    # ------------------------------------------------------------------

    async def _build_resource_plan(
        self,
        plan: ExecutionPlan,
        step_contexts: list[StepContext],
    ) -> ResourcePlan:
        """Single LLM call to synthesize gathered contexts into execution briefs."""
        # Serialise gathered contexts for the LLM
        contexts_payload = []
        for sc in step_contexts:
            contexts_payload.append({
                "step_number": sc.step_number,
                "step_description": sc.step_description,
                "gathered_contexts": [
                    {
                        "source": c.source.value,
                        "query": c.query,
                        "content": c.content[:800] if c.content else "",
                        "source_detail": c.source_detail,
                    }
                    for c in sc.gathered_contexts
                    if c.content  # skip empty placeholders (step_dep)
                ],
            })

        import json as _json
        user_msg = (
            f"Plan objective: {plan.objective}\n\n"
            f"Gathered contexts per step:\n{_json.dumps(contexts_payload, indent=2)}\n\n"
            "Synthesize these into execution-ready briefs per step."
        )

        kb_docs_used: list[str] = []
        web_count = 0
        for sc in step_contexts:
            for c in sc.gathered_contexts:
                if c.source == ContextSource.KNOWLEDGE_BASE and c.source_detail:
                    if c.source_detail not in kb_docs_used:
                        kb_docs_used.append(c.source_detail)
                elif c.source == ContextSource.WEB_SEARCH:
                    web_count += 1

        try:
            result = await self.llm.generate_json(CONTEXT_SUMMARY_PROMPT, user_msg)
        except Exception:
            logger.exception("ContextGatheringAgent: resource plan summary LLM call failed")
            result = {}

        # Enrich step contexts with summaries from LLM
        briefs_map: dict[int, dict] = {
            b["step_number"]: b
            for b in result.get("step_briefs", [])
            if isinstance(b, dict) and "step_number" in b
        }
        for sc in step_contexts:
            brief = briefs_map.get(sc.step_number, {})
            sc.context_summary = brief.get("context_summary", "")
            sc.execution_strategy = brief.get("execution_strategy", "")

        # Build dependency graph: step_number -> list of dependency step numbers
        dep_graph: dict[str, list[int]] = {}
        for s in plan.steps:
            dep_graph[str(s.step_number)] = s.dependencies

        # Also check for explicit step_dep requirements
        result_dep_graph = result.get("dependency_graph", dep_graph)

        return ResourcePlan(
            objective=plan.objective,
            total_steps=len(plan.steps),
            tool_inventory=result.get("tool_inventory", []),
            step_contexts=step_contexts,
            dependency_graph=result_dep_graph,
            knowledge_base_docs_used=kb_docs_used,
            web_searches_performed=web_count,
            estimated_complexity=result.get("estimated_complexity", "medium"),
        )