"""Dynamic plan updater for adaptive execution.

After each step completes, this module evaluates the results and updates
the plan accordingly - adding new steps if needed or removing steps that
are no longer necessary.
"""

from __future__ import annotations

from uuid import uuid4
from typing import Any

from promethian.core.types import Plan, PlanStep, StepResult, ExecutionStatus
from promethian.utils.claude import ClaudeClient


PLAN_UPDATE_SYSTEM = """You are a plan adaptation assistant for an AI agent.

After a step completes, you evaluate whether the remaining plan needs adjustment.

Your job is to:
1. Analyze what the completed step achieved
2. Determine if remaining steps are still needed or should be modified
3. Identify if new steps are required based on the results
4. Provide the updated plan

Be conservative - only add steps when truly necessary and remove steps only
when they are clearly no longer needed. Avoid unnecessary plan churn.

Output your decision as JSON."""


class PlanUpdater:
    """
    Updates execution plans based on step results.

    After each step, evaluates whether the plan should be modified:
    - Add new steps if results reveal more work needed
    - Remove steps that are no longer necessary
    - Modify remaining steps based on new information
    """

    def __init__(self, claude: ClaudeClient):
        self.claude = claude

    async def should_update_plan(
        self,
        plan: Plan,
        completed_step: PlanStep,
        step_result: StepResult,
        remaining_steps: list[PlanStep],
    ) -> bool:
        """
        Quick check if a plan update is warranted.

        Returns True if we should run the full plan update.
        Avoids expensive LLM calls when not needed.
        """
        # Always update if step failed - may need to add recovery steps
        if step_result.status != ExecutionStatus.SUCCESS:
            return True

        # No remaining steps - nothing to update
        if not remaining_steps:
            return False

        # If output contains certain signals, update
        output_str = str(step_result.output).lower() if step_result.output else ""
        update_signals = [
            "need to also",
            "additionally",
            "discovered that",
            "found that we need",
            "unexpectedly",
            "however",
            "but first",
            "before we can",
            "prerequisite",
            "dependency",
            "error",
            "failed",
            "couldn't",
            "unable to",
        ]

        if any(signal in output_str for signal in update_signals):
            return True

        # If tool results contain errors or unexpected data
        if isinstance(step_result.output, dict):
            tool_results = step_result.output.get("tool_results", [])
            for result in tool_results:
                if not result.get("success", True):
                    return True
                # Check if result suggests more work
                result_str = str(result.get("result", "")).lower()
                if any(signal in result_str for signal in update_signals):
                    return True

        return False

    async def update_plan(
        self,
        plan: Plan,
        completed_step: PlanStep,
        step_result: StepResult,
        step_outputs: dict[str, Any],
        remaining_steps: list[PlanStep],
        user_feedback: str | None = None,
    ) -> tuple[list[PlanStep], str]:
        """
        Update the plan based on step results and/or user feedback.

        Args:
            plan: The original plan
            completed_step: The step that just completed
            step_result: Result of the completed step
            step_outputs: All step outputs so far
            remaining_steps: Steps that haven't been executed yet
            user_feedback: Optional real-time feedback from the user

        Returns:
            Tuple of (updated_remaining_steps, reasoning)
        """
        # Format the step result for the prompt
        result_summary = self._format_step_result(step_result)

        # Format remaining steps
        remaining_summary = self._format_remaining_steps(remaining_steps)

        # Build the prompt based on whether we have user feedback
        has_feedback = user_feedback is not None
        prompt_intro = "A plan step has completed"
        if has_feedback:
            prompt_intro += " and the user has provided feedback that must be incorporated"
        prompt_intro += ". Evaluate if the remaining plan needs updating."

        prompt = f"""{prompt_intro}

## Original Task
{plan.task}

## Completed Step
ID: {completed_step.id}
Description: {completed_step.description}
Expected Output: {completed_step.expected_output}

## Step Result
Status: {step_result.status.value}
{result_summary}

## Remaining Planned Steps
{remaining_summary if remaining_summary else "No remaining steps planned."}

## Available Resources
Skills: {[s.name for s in plan.relevant_skills]}
Tools: {[t.name for t in plan.relevant_tools]}
{f'''
## USER FEEDBACK (IMPORTANT - Must Address!)
The user has provided the following real-time feedback that should be incorporated:

{user_feedback}

This feedback takes priority. Adjust the plan to address the user's concerns or direction.
''' if user_feedback else ''}
## Decision Required

Based on the step result{' and user feedback' if user_feedback else ''}, determine:
1. Are any remaining steps no longer needed? (e.g., already accomplished, not relevant)
2. Do any remaining steps need modification based on what we learned?
3. Are new steps needed that weren't in the original plan?

Be conservative:
- Only remove steps if they're clearly unnecessary
- Only add steps if the result revealed genuinely new requirements
- Keep modifications minimal

Respond with JSON:
{{
    "needs_update": true/false,
    "reasoning": "Brief explanation of why changes are needed (or not)",
    "steps_to_remove": ["step_id_1"],  // IDs of steps to remove
    "steps_to_modify": [
        {{
            "id": "existing_step_id",
            "new_description": "Updated description",
            "new_expected_output": "Updated expected output"
        }}
    ],
    "steps_to_add": [
        {{
            "description": "What this new step should do",
            "expected_output": "What it should produce",
            "insert_position": "next" | "end" | "before:step_id",
            "skills_to_use": [],
            "tools_to_use": [],
            "tools_to_build": []
        }}
    ]
}}"""

        try:
            result = await self.claude.complete_json(
                prompt=prompt,
                system=PLAN_UPDATE_SYSTEM,
            )
        except Exception as e:
            # On failure, return unchanged plan
            return remaining_steps, f"Plan update failed: {e}"

        if not result.get("needs_update", False):
            return remaining_steps, result.get("reasoning", "No changes needed")

        # Apply updates
        updated_steps = self._apply_updates(
            remaining_steps=remaining_steps,
            steps_to_remove=result.get("steps_to_remove", []),
            steps_to_modify=result.get("steps_to_modify", []),
            steps_to_add=result.get("steps_to_add", []),
        )

        return updated_steps, result.get("reasoning", "Plan updated")

    def _format_step_result(self, step_result: StepResult) -> str:
        """Format a step result for the prompt."""
        parts = []

        if step_result.output:
            if isinstance(step_result.output, dict):
                # Handle structured output with tool results
                if "text" in step_result.output:
                    parts.append(f"Response: {step_result.output['text'][:500]}")
                if "tool_results" in step_result.output:
                    parts.append("\nTool Results:")
                    for i, tr in enumerate(step_result.output["tool_results"][:5]):
                        status = "✓" if tr.get("success") else "✗"
                        result_preview = str(tr.get("result", tr.get("error", "")))[:200]
                        parts.append(f"  {status} Tool {i+1}: {result_preview}")
            else:
                parts.append(f"Output: {str(step_result.output)[:500]}")

        if step_result.error:
            parts.append(f"Error: {step_result.error}")

        if step_result.healing_log:
            parts.append(f"Healing attempts: {len(step_result.healing_log)}")

        return "\n".join(parts) if parts else "No output captured"

    def _format_remaining_steps(self, remaining_steps: list[PlanStep]) -> str:
        """Format remaining steps for the prompt."""
        if not remaining_steps:
            return ""

        parts = []
        for step in remaining_steps:
            deps = f" (depends on: {', '.join(step.dependencies)})" if step.dependencies else ""
            parts.append(f"- [{step.id}] {step.description}{deps}")
            parts.append(f"  Expected: {step.expected_output}")

        return "\n".join(parts)

    def _apply_updates(
        self,
        remaining_steps: list[PlanStep],
        steps_to_remove: list[str],
        steps_to_modify: list[dict],
        steps_to_add: list[dict],
    ) -> list[PlanStep]:
        """Apply the updates to the remaining steps."""
        # Create a mutable copy
        updated = list(remaining_steps)

        # Remove steps
        removed_set = set(steps_to_remove)
        updated = [s for s in updated if s.id not in removed_set]

        # IMPORTANT: Remove references to deleted steps from dependencies
        # Otherwise remaining steps will fail with "dependency not satisfied"
        for i, step in enumerate(updated):
            if step.dependencies:
                new_deps = [d for d in step.dependencies if d not in removed_set]
                if new_deps != step.dependencies:
                    # Create new step with updated dependencies
                    updated[i] = PlanStep(
                        id=step.id,
                        description=step.description,
                        expected_output=step.expected_output,
                        skills_to_use=step.skills_to_use,
                        tools_to_use=step.tools_to_use,
                        tools_to_build=step.tools_to_build,
                        dependencies=new_deps,
                    )

        # Modify steps
        modify_map = {m["id"]: m for m in steps_to_modify}
        for i, step in enumerate(updated):
            if step.id in modify_map:
                mod = modify_map[step.id]
                updated[i] = PlanStep(
                    id=step.id,
                    description=mod.get("new_description", step.description),
                    expected_output=mod.get("new_expected_output", step.expected_output),
                    skills_to_use=mod.get("skills_to_use", step.skills_to_use),
                    tools_to_use=mod.get("tools_to_use", step.tools_to_use),
                    tools_to_build=mod.get("tools_to_build", step.tools_to_build),
                    dependencies=mod.get("dependencies", step.dependencies),
                )

        # Add new steps
        for new_step in steps_to_add:
            step = PlanStep(
                id=f"step_{uuid4().hex[:8]}",
                description=new_step.get("description", ""),
                expected_output=new_step.get("expected_output", ""),
                skills_to_use=new_step.get("skills_to_use", []),
                tools_to_use=new_step.get("tools_to_use", []),
                tools_to_build=new_step.get("tools_to_build", []),
                dependencies=new_step.get("dependencies", []),
            )

            position = new_step.get("insert_position", "end")
            if position == "next":
                updated.insert(0, step)
            elif position == "end":
                updated.append(step)
            elif position.startswith("before:"):
                target_id = position[7:]
                idx = next((i for i, s in enumerate(updated) if s.id == target_id), len(updated))
                updated.insert(idx, step)
            else:
                updated.append(step)

        return updated


async def evaluate_tool_results_for_next_action(
    claude: ClaudeClient,
    step_description: str,
    tool_results: list[dict],
    available_tools: list[dict],
) -> tuple[bool, list[dict] | None, str | None]:
    """
    Evaluate tool results to decide if more tool calls are needed.

    This is the core of the agentic loop within a single step.

    Args:
        claude: Claude client
        step_description: What the step is trying to accomplish
        tool_results: Results from the previous tool calls
        available_tools: Tools that can be called

    Returns:
        Tuple of (should_continue, next_tool_calls, final_response)
        - If should_continue=True, next_tool_calls contains the tools to call
        - If should_continue=False, final_response contains the step output
    """
    # Format tool results
    results_summary = []
    for i, result in enumerate(tool_results):
        if result.get("success"):
            results_summary.append(f"Tool {i+1}: SUCCESS\n{str(result.get('result', ''))[:500]}")
        else:
            results_summary.append(f"Tool {i+1}: FAILED\n{result.get('error', 'Unknown error')}")

    prompt = f"""You are executing a step that may require multiple tool calls.

## Step Goal
{step_description}

## Tool Results from Previous Call(s)
{chr(10).join(results_summary)}

## Available Tools
{[t['name'] for t in available_tools]}

Based on these results, decide:
1. Have we accomplished the step goal?
2. Do we need to call more tools to complete the step?
3. If calling more tools, which ones and with what parameters?

Respond with JSON:
{{
    "step_complete": true/false,
    "reasoning": "Why the step is complete or what more is needed",
    "next_tool_calls": [
        {{
            "name": "tool_name",
            "input": {{...}}
        }}
    ],
    "final_response": "Summary of what was accomplished (only if step_complete=true)"
}}"""

    try:
        result = await claude.complete_json(prompt=prompt)
    except Exception:
        # On failure, consider step complete with current results
        return False, None, f"Completed with results: {results_summary}"

    if result.get("step_complete", True):
        return False, None, result.get("final_response", "Step completed")

    next_calls = result.get("next_tool_calls", [])
    if not next_calls:
        return False, None, result.get("final_response", "Step completed")

    return True, next_calls, None
