execution_agent.py

from __future__ import annotations

import json
import logging
import re
from typing import Any, Awaitable, Callable, Optional

from backend.agents.base import BaseAgent
from backend.models.enums import AgentPhase
from backend.models.schemas import (
    ExecutionPlan,
    ExecutionResult,
    PlanStep,
    ResourcePlan,
    StepContext,
    StepResult,
)
from backend.prompts.context_prompt import FINAL_ASSEMBLY_PROMPT, STEP_EXECUTION_PROMPT

logger = logging.getLogger(__name__)


class ExecutionAgent(BaseAgent):
    phase = AgentPhase.EXECUTION
    name = "Execution Agent"

    def get_system_prompt(self) -> str:
        return STEP_EXECUTION_PROMPT

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

    async def process(
        self,
        input_data: Any,
        on_stream: Optional[Callable[[str], Awaitable[None]]] = None,
        on_step_update: Optional[Callable[[int, str, dict], Awaitable[None]]] = None,
    ) -> ExecutionResult:
        """Execute a plan step-by-step, using the ResourcePlan for enriched context.

        When a ResourcePlan is provided (Phase 4.5 output), each step is executed
        independently with its gathered context injected.  Falls back to the
        original single-call approach when no ResourcePlan is available.

        Args:
            input_data:     ExecutionPlan, or a dict with keys 'plan' and
                            optionally 'resource_plan'.
            on_stream:      Called with each streamed text chunk.
            on_step_update: Called as (step_number, status, data) for UI updates.
        """
        if isinstance(input_data, dict):
            plan: ExecutionPlan = input_data["plan"]
            resource_plan: Optional[ResourcePlan] = input_data.get("resource_plan")
        else:
            plan = input_data
            resource_plan = None

        if resource_plan:
            return await self._execute_step_by_step(
                plan, resource_plan, on_stream, on_step_update
            )
        return await self._execute_single_call(plan, on_stream)

    async def process_step(
        self,
        step: PlanStep,
        step_context: Optional[StepContext],
        prior_step_outputs: dict[int, StepResult],
        on_stream: Optional[Callable[[str], Awaitable[None]]] = None,
    ) -> StepResult:
        """Execute a single plan step with gathered context and prior outputs.

        Args:
            step:               The PlanStep to execute.
            step_context:       Pre-gathered context from the ResourcePlan.
            prior_step_outputs: Results of already-completed steps (keyed by step number).
            on_stream:          Streaming callback for UI.
        """
        user_msg = self._build_step_user_message(step, step_context, prior_step_outputs)

        if on_stream:
            raw = await self.llm.generate_streaming(
                STEP_EXECUTION_PROMPT, user_msg, on_stream
            )
        else:
            raw = await self.llm.generate(STEP_EXECUTION_PROMPT, user_msg)

        parsed = self._parse_response(raw)
        return StepResult(
            step_number=step.step_number,
            output=parsed.get("output", ""),
            status=parsed.get("status", "completed"),
            validation_passed=parsed.get("validation_passed", True),
            contexts_used=parsed.get("contexts_used", []),
            reasoning=parsed.get("reasoning", ""),
        )

    # ------------------------------------------------------------------
    # Internal: step-by-step execution
    # ------------------------------------------------------------------

    async def _execute_step_by_step(
        self,
        plan: ExecutionPlan,
        resource_plan: ResourcePlan,
        on_stream: Optional[Callable[[str], Awaitable[None]]],
        on_step_update: Optional[Callable[[int, str, dict], Awaitable[None]]],
    ) -> ExecutionResult:
        """Iterate through steps in dependency order, executing each individually."""
        # Build a map from step_number -> StepContext
        context_map: dict[int, StepContext] = {
            sc.step_number: sc for sc in resource_plan.step_contexts
        }

        # Topological ordering respecting dependencies
        ordered_steps = self._topological_sort(plan.steps)

        step_results: dict[int, StepResult] = {}

        for step in ordered_steps:
            step_ctx = context_map.get(step.step_number)

            if on_step_update:
                await on_step_update(
                    step.step_number,
                    "started",
                    {
                        "description": step.description,
                        "total_steps": len(ordered_steps),
                        "context_sources": (
                            [c.source.value for c in step_ctx.gathered_contexts]
                            if step_ctx else []
                        ),
                    },
                )

            # Create a streaming wrapper that prefixes each chunk with step info
            async def step_stream(chunk: str, sn: int = step.step_number) -> None:
                if on_stream:
                    await on_stream(chunk)

            try:
                result = await self.process_step(
                    step,
                    step_ctx,
                    step_results,
                    on_stream=step_stream,
                )
            except Exception as exc:
                logger.exception("Step %d execution failed", step.step_number)
                result = StepResult(
                    step_number=step.step_number,
                    output=f"Step failed: {exc}",
                    status="failed",
                    validation_passed=False,
                )

            step_results[step.step_number] = result

            if on_step_update:
                await on_step_update(
                    step.step_number,
                    "completed",
                    {
                        "status": result.status,
                        "validation_passed": result.validation_passed,
                        "output_preview": result.output[:200] if result.output else "",
                    },
                )

        # Final assembly: quality checks + trace_to_goal
        all_results = [step_results[s.step_number] for s in plan.steps if s.step_number in step_results]
        final = await self._assemble_final(plan, all_results, on_stream)

        return ExecutionResult(
            plan=plan,
            resource_plan=resource_plan,
            step_results=all_results,
            final_output=final.get("final_output", ""),
            completeness_check=final.get("completeness_check", ""),
            clarity_check=final.get("clarity_check", ""),
            relevance_check=final.get("relevance_check", ""),
            correctness_check=final.get("correctness_check", ""),
            trace_to_goal=final.get("trace_to_goal", ""),
        )

    async def _assemble_final(
        self,
        plan: ExecutionPlan,
        step_results: list[StepResult],
        on_stream: Optional[Callable[[str], Awaitable[None]]],
    ) -> dict:
        """One final LLM call to produce quality checks and the complete final output."""
        steps_summary = "\n".join(
            f"Step {r.step_number} ({r.status}): {r.output[:500]}"
            for r in step_results
        )
        user_msg = (
            f"Plan objective: {plan.objective}\n\n"
            f"Plan deliverables: {plan.deliverables}\n\n"
            f"Completed step outputs:\n{steps_summary}\n\n"
            "Assemble the final deliverable and run quality checks."
        )
        if on_stream:
            raw = await self.llm.generate_streaming(FINAL_ASSEMBLY_PROMPT, user_msg, on_stream)
        else:
            raw = await self.llm.generate(FINAL_ASSEMBLY_PROMPT, user_msg)

        try:
            return self._parse_response(raw)
        except Exception:
            logger.exception("Final assembly parsing failed")
            return {
                "final_output": "\n\n".join(r.output for r in step_results if r.output),
                "completeness_check": "Auto-assembled from step outputs",
                "trace_to_goal": f"Steps executed for: {plan.objective}",
            }

    # ------------------------------------------------------------------
    # Internal: legacy single-call fallback (when no ResourcePlan)
    # ------------------------------------------------------------------

    async def _execute_single_call(
        self,
        plan: ExecutionPlan,
        on_stream: Optional[Callable[[str], Awaitable[None]]],
    ) -> ExecutionResult:
        """Fallback: send the entire plan in one call (original behaviour)."""
        from backend.prompts.execution_prompt import EXECUTION_SYSTEM_PROMPT  # type: ignore[import]

        user_msg = (
            "Execute the following plan step by step.\n\n"
            f"Plan:\n{plan.model_dump_json(indent=2)}\n\n"
            "For each step: produce output, validate it, report status.\n"
            "After all steps: produce the final deliverable and run "
            "completeness/clarity/relevance/correctness checks.\n"
            "Include an explicit trace from the output back to the original goal.\n\n"
            "Do NOT include the plan object in your response.\n"
            "Respond ONLY with JSON."
        )
        if on_stream:
            raw = await self.llm.generate_streaming(EXECUTION_SYSTEM_PROMPT, user_msg, on_stream)
        else:
            raw = await self.llm.generate(EXECUTION_SYSTEM_PROMPT, user_msg)

        parsed = self._parse_response(raw)
        parsed["plan"] = plan.model_dump()
        return ExecutionResult(**parsed)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _build_step_user_message(
        self,
        step: PlanStep,
        step_context: Optional[StepContext],
        prior_step_outputs: dict[int, StepResult],
    ) -> str:
        parts = [
            f"## Step {step.step_number}: {step.description}",
            f"Expected output: {step.expected_output}" if step.expected_output else "",
            f"Validation criteria: {step.validation}" if step.validation else "",
        ]

        if step_context:
            if step_context.context_summary:
                parts.append(
                    f"\n## Context gathered for this step\n{step_context.context_summary}"
                )
            if step_context.execution_strategy:
                parts.append(
                    f"\n## Suggested execution strategy\n{step_context.execution_strategy}"
                )
            kb_contexts = [
                c for c in step_context.gathered_contexts
                if c.source.value == "knowledge_base" and c.content
            ]
            if kb_contexts:
                kb_section = "\n\n".join(
                    f"[{c.source_detail}]\n{c.content}"
                    for c in kb_contexts
                )
                parts.append(f"\n## Knowledge base findings\n{kb_section}")
            web_contexts = [
                c for c in step_context.gathered_contexts
                if c.source.value == "web_search" and c.content
            ]
            if web_contexts:
                web_section = "\n\n".join(
                    f"[{c.metadata.get('title', c.source_detail)}]({c.source_detail})\n{c.content}"
                    for c in web_contexts
                )
                parts.append(f"\n## Web search findings\n{web_section}")

        # Inject outputs from dependency steps
        if step.dependencies:
            dep_outputs = {
                dep: prior_step_outputs[dep]
                for dep in step.dependencies
                if dep in prior_step_outputs
            }
            if dep_outputs:
                dep_section = "\n\n".join(
                    f"Step {dep_num} output:\n{res.output}"
                    for dep_num, res in dep_outputs.items()
                )
                parts.append(f"\n## Outputs from prerequisite steps\n{dep_section}")

        parts.append("\nExecute this step and produce the JSON response.")
        return "\n".join(p for p in parts if p)

    @staticmethod
    def _topological_sort(steps: list[PlanStep]) -> list[PlanStep]:
        """Return steps in an order that respects their dependency constraints."""
        step_map = {s.step_number: s for s in steps}
        visited: set[int] = set()
        ordered: list[PlanStep] = []

        def visit(num: int) -> None:
            if num in visited:
                return
            visited.add(num)
            step = step_map.get(num)
            if step:
                for dep in step.dependencies:
                    visit(dep)
                ordered.append(step)

        for s in steps:
            visit(s.step_number)
        return ordered

    @staticmethod
    def _parse_response(text: str) -> dict:
        cleaned = text.strip()
        cleaned = re.sub(r"<think>.*?</think>", "", cleaned, flags=re.DOTALL).strip()
        if cleaned.startswith("```"):
            cleaned = re.sub(r"^```\w*\n?", "", cleaned)
            cleaned = re.sub(r"\n?```$", "", cleaned)
        match = re.search(r"\{[\s\S]*\}", cleaned)
        if match:
            return json.loads(match.group())
        raise ValueError(f"Could not parse execution result JSON: {text[:200]}")