from __future__ import annotations

import ast
import re
from typing import Any
from pathlib import Path

from langgraph.graph import END, START, StateGraph
from typing_extensions import TypedDict

from ai_orchestrator.config.settings import Settings
from ai_orchestrator.orchestrators.base import OrchestratorEngine
from ai_orchestrator.schemas.task_request import TaskRequest
from ai_orchestrator.schemas.task_state import (
    ArtifactRecord,
    CommandRecord,
    StepRecord,
    TaskState,
    TestResult,
)
from ai_orchestrator.schemas.task_status import TaskLifecycleStatus
from ai_orchestrator.schemas.tool_result import ToolResult
from ai_orchestrator.schemas.workflow_models import CriticResult, StructuredError, ValidationSummary
from ai_orchestrator.services.agent_controller import AgentController
from ai_orchestrator.services.artifact_store import ArtifactStore
from ai_orchestrator.services.embedding_service import EmbeddingService
from ai_orchestrator.services.model_router import ModelRouter
from ai_orchestrator.services.prompt_loader import PromptLoader
from ai_orchestrator.services.rag_knowledge_base import RAGKnowledgeBase
from ai_orchestrator.services.sqlite_store import SQLiteTaskStore
from ai_orchestrator.tools.file_tools import FileTools
from ai_orchestrator.tools.git_runner import GitRunner


class WorkflowState(TypedDict, total=False):
    task_id: str
    goal: str
    source_repo_path: str
    max_attempts: int
    attempt: int
    test_command: str | None
    repo_context: dict[str, Any]
    change_mode: str
    patch_text: str
    patch_rationale: str
    file_writes: list[dict[str, str]]
    critic_result: dict[str, Any]
    apply_result: dict[str, Any]
    validation_result: dict[str, Any]
    evaluation: dict[str, Any]
    failure_context: str | None
    final_diff: str
    summary: str


class LangGraphEngine(OrchestratorEngine):
    def __init__(
        self,
        *,
        settings: Settings,
        store: SQLiteTaskStore,
        artifact_store: ArtifactStore,
        prompt_loader: PromptLoader,
        model_router: ModelRouter,
        file_tools: FileTools,
        git_runner: GitRunner,
        agent_controller: AgentController,
        rag_knowledge_base: RAGKnowledgeBase | None = None,
        embedding_service: EmbeddingService | None = None,
    ) -> None:
        self.settings = settings
        self.store = store
        self.artifact_store = artifact_store
        self.prompt_loader = prompt_loader
        self.model_router = model_router
        self.file_tools = file_tools
        self.git_runner = git_runner
        self.agent_controller = agent_controller
        self.rag_knowledge_base = rag_knowledge_base
        self.embedding_service = embedding_service
        self.graph = self._build_graph()

    def _build_graph(self):
        builder = StateGraph(WorkflowState)
        builder.add_node("create_snapshot", self._create_snapshot)
        builder.add_node("inspect_repo", self._inspect_repo)
        builder.add_node("plan_task", self._plan_task)
        builder.add_node("generate_patch", self._generate_patch)
        builder.add_node("critic_patch", self._critic_patch)
        builder.add_node("apply_patch", self._apply_patch)
        builder.add_node("run_validation", self._run_validation)
        builder.add_node("evaluate_result", self._evaluate_result)
        builder.add_node("summarize", self._summarize)

        builder.add_edge(START, "create_snapshot")
        builder.add_edge("create_snapshot", "inspect_repo")
        builder.add_edge("inspect_repo", "plan_task")
        builder.add_edge("plan_task", "generate_patch")
        builder.add_edge("generate_patch", "critic_patch")
        builder.add_conditional_edges(
            "critic_patch",
            self._route_after_critic,
            {"retry": "generate_patch", "apply": "apply_patch"},
        )
        builder.add_edge("apply_patch", "run_validation")
        builder.add_edge("run_validation", "evaluate_result")
        builder.add_conditional_edges(
            "evaluate_result",
            self._route_after_evaluation,
            {"retry": "generate_patch", "summarize": "summarize"},
        )
        builder.add_edge("summarize", END)
        return builder.compile()

    async def submit_task(self, task_id: str, request: TaskRequest) -> None:
        task = self._require_task(task_id)
        initial_state: WorkflowState = {
            "task_id": task_id,
            "goal": request.goal,
            "source_repo_path": str(request.repo_path),
            "max_attempts": request.constraints.max_attempts,
            "attempt": task.attempt,
            "test_command": request.constraints.test_command,
            "failure_context": None,
        }
        await self.graph.ainvoke(initial_state)

    async def get_status(self, task_id: str):
        return self.store.get(task_id)

    async def get_result(self, task_id: str):
        return self.store.get(task_id)

    async def cancel_task(self, task_id: str) -> None:
        task = self.store.get(task_id)
        if task is None or task.status in TaskLifecycleStatus.terminal():
            return
        task.mark_terminal(TaskLifecycleStatus.cancelled, summary="Task cancelled.")
        self.store.save(task)

    async def _create_snapshot(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("create_snapshot")
        self.store.save(task)

        await self.agent_controller.create_workspace(task)
        self._add_artifact(
            task,
            self.artifact_store.write_json(
                task.task_id,
                "workspace_snapshot.json",
                {
                    "source_repo_path": str(task.repo_path),
                    "workspace_path": str(task.workspace_path),
                    "strategy": task.workspace_strategy,
                },
            ),
        )
        self._append_step(task, "create_snapshot", "Executor", "completed", task.workspace_strategy)
        self.store.save(task)
        return {}

    async def _inspect_repo(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("inspect_repo")
        self.store.save(task)

        repo_context = await self.agent_controller.analyze_repository(task)
        focus_issue_inventory = await self.model_router.identify_focus_issues(
            task=task,
            repo_snapshot=repo_context,
            instruction=self.prompt_loader.load("issue_analyzer.txt"),
        )
        if focus_issue_inventory.focus_paths:
            task.focus_issue_inventory = focus_issue_inventory
            repo_context["focus_issue_inventory"] = focus_issue_inventory.model_dump(mode="json")
            repo_context["acceptance_checks"] = focus_issue_inventory.acceptance_checks
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    "focus_issue_inventory.json",
                    focus_issue_inventory.model_dump(mode="json"),
                ),
            )
        self._add_artifact(
            task,
            self.artifact_store.write_json(task.task_id, "repo_snapshot.json", repo_context),
        )
        if task.memory_hits:
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    "memory_hits.json",
                    [hit.model_dump(mode="json") for hit in task.memory_hits],
                ),
            )
        self._append_step(task, "inspect_repo", "Planner", "completed", f"Indexed {len(repo_context.get('key_modules', []))} key modules.")
        self._advance_subtask(task)
        self.store.save(task)
        return {"repo_context": repo_context}

    async def _plan_task(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("plan_task")
        self.store.save(task)

        plan_steps = await self.agent_controller.plan_task(
            task,
            instruction=self.prompt_loader.load("planner.txt"),
        )
        self._add_artifact(
            task,
            self.artifact_store.write_json(
                task.task_id,
                "plan.json",
                {"goal": task.goal, "steps": plan_steps},
            ),
        )
        self._add_artifact(
            task,
            self.artifact_store.write_json(
                task.task_id,
                "subtasks.json",
                [subtask.model_dump(mode="json") for subtask in task.subtasks],
            ),
        )
        self._append_step(task, "plan_task", "Planner", "completed", "Planner and supervisor prepared execution subtasks.")
        self._advance_subtask(task)
        self.store.save(task)
        return {"repo_context": task.repo_context}

    async def _generate_patch(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("generate_patch")
        task.attempt += 1
        self.store.save(task)

        diff_before_result = await self.git_runner.diff_repo(
            repo_path=task.execution_repo_path(),
            base_path=task.repo_path if task.workspace_strategy == "copy" else None,
        )
        diff_before = diff_before_result.stdout if diff_before_result.success else ""
        patterns = []
        if state.get("failure_context"):
            patterns = await self.agent_controller.memory_service.retrieve_patterns(
                state["failure_context"],
                overrides=task.constraints.model_overrides,
            )

        repo_snapshot = dict(task.repo_context)
        if task.memory_hits:
            repo_snapshot["memory_hits"] = [hit.model_dump(mode="json") for hit in task.memory_hits]
        if patterns:
            repo_snapshot["pattern_memory"] = patterns
        retry_focus_files = list(task.retry_focus_files or repo_snapshot.get("failure_focus_files", []))
        if retry_focus_files:
            expanded_retry_focus = self._expand_repo_snapshot(
                task=task,
                repo_snapshot=repo_snapshot,
                requested_paths=retry_focus_files,
            )
            if expanded_retry_focus:
                repo_snapshot["failure_focus_files"] = expanded_retry_focus

        # RAG: Retrieve similar bug-fix examples from knowledge base
        rag_examples = []
        if self.rag_knowledge_base and self.embedding_service:
            rag_examples = await self._retrieve_rag_examples(task, repo_snapshot)
        if rag_examples:
            repo_snapshot["rag_examples"] = rag_examples

        # Progressive retry strategy: different approaches for each attempt
        attempt_instruction = self._get_attempt_instruction(task.attempt, task.constraints.max_attempts)
        repo_snapshot["attempt_number"] = task.attempt
        repo_snapshot["total_attempts"] = task.constraints.max_attempts

        proposal = None
        for context_round in range(self.settings.orchestrator_context_request_rounds):
            proposal = await self.model_router.generate_patch(
                task=task,
                repo_snapshot=repo_snapshot,
                diff_before=diff_before,
                failure_context=state.get("failure_context"),
                instruction=attempt_instruction,
            )
            if not proposal.context_requests:
                break

            expanded = self._expand_repo_snapshot(
                task=task,
                repo_snapshot=repo_snapshot,
                requested_paths=[request.path for request in proposal.context_requests],
            )
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    f"attempt_{task.attempt}_context_request_{context_round + 1}.json",
                    {
                        "requests": [
                            {"path": request.path, "reason": request.reason}
                            for request in proposal.context_requests
                        ],
                        "expanded_files": expanded,
                    },
                ),
            )
            if not expanded:
                break
            task.repo_context = repo_snapshot
            self.store.save(task)
        assert proposal is not None

        self._add_artifact(
            task,
            self.artifact_store.write_text(
                task.task_id,
                f"attempt_{task.attempt}_coder_notes.txt",
                proposal.rationale,
            ),
        )
        if proposal.patch.strip():
            self._add_artifact(
                task,
                self.artifact_store.write_text(
                    task.task_id,
                    f"attempt_{task.attempt}_patch.diff",
                    proposal.patch,
                ),
            )
        if proposal.file_writes:
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    f"attempt_{task.attempt}_file_writes.json",
                    [
                        {"path": file_write.path, "content": file_write.content}
                        for file_write in proposal.file_writes
                    ],
                ),
            )
        self._append_step(task, "generate_patch", "Coder", "completed", proposal.rationale)
        self._advance_subtask(task)
        self.store.save(task)
        return {
            "attempt": task.attempt,
            "change_mode": proposal.mode,
            "patch_text": proposal.patch,
            "patch_rationale": proposal.rationale,
            "file_writes": [
                {"path": file_write.path, "content": file_write.content}
                for file_write in proposal.file_writes
            ],
        }

    async def _critic_patch(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("critic_patch")
        self.store.save(task)

        # Ensemble review: multiple reviewers vote on patch quality
        num_reviewers = self.settings.critic_reviewer_count
        votes = []
        critic_results = []  # Store full results to access issues
        reviewer_prompts = [
            "You are Reviewer 1. Be strict about correctness.",
            "You are Reviewer 2. Focus on edge cases and error handling.",
            "You are Reviewer 3. Check for security issues and best practices.",
        ]
        while len(reviewer_prompts) < num_reviewers:
            reviewer_prompts.append(f"You are Reviewer {len(reviewer_prompts) + 1}. Focus on correctness and completeness.")

        for i in range(num_reviewers):
            # Customize the critic prompt for this reviewer
            original_prompt = self.prompt_loader.load("critic.txt")
            reviewer_prompt = f"{original_prompt}\n\n{reviewer_prompts[i]}"

            critic_result = self.agent_controller.reflect_patch(
                task,
                state.get("patch_text", ""),
                file_writes=state.get("file_writes", []),
                instruction=reviewer_prompt,
            )

            # Store individual reviewer result
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    f"attempt_{task.attempt}_critic_{i+1}.json",
                    {
                        **critic_result.model_dump(mode="json"),
                        "reviewer": i + 1,
                        "prompt_addition": reviewer_prompts[i],
                    },
                ),
            )

            votes.append(critic_result.result)
            critic_results.append(critic_result)

        # Ensemble decision: require consensus to proceed
        approve_count = sum(1 for vote in votes if vote == "pass")
        reject_count = sum(1 for vote in votes if vote == "fail")

        # Final decision: pass only if majority approve, reject if any critical issue found
        final_result = "pass" if approve_count >= (num_reviewers // 2 + 1) else "fail"

        # Collect all issues from all reviewers
        all_issues = []
        for i, critic_result in enumerate(critic_results):
            if hasattr(critic_result, 'issues') and critic_result.issues:
                all_issues.extend([f"Reviewer{i+1}: {issue}" for issue in critic_result.issues])

        ensemble_result = type('obj', (object,), {
            'result': final_result,
            'issues': all_issues,
            'vote_breakdown': f"{approve_count} approve, {reject_count} reject",
            'ensemble_size': num_reviewers,
        })

        self._add_artifact(
            task,
            self.artifact_store.write_json(
                task.task_id,
                f"attempt_{task.attempt}_critic_ensemble.json",
                {
                    "result": final_result,
                    "vote_breakdown": f"{approve_count} approve, {reject_count} reject",
                    "individual_votes": [v for v in votes],
                    "all_issues": all_issues,
                    "ensemble_size": num_reviewers,
                },
            ),
        )

        details = f"Ensemble review ({approve_count}/{num_reviewers} approve): "
        if final_result == "pass":
            details += "PASSED"
        else:
            details += f"REJECTED - {', '.join(all_issues[:3])}" if all_issues else "REJECTED"

        self._append_step(task, "critic_patch", "EnsembleCritic", "completed", details)
        self.store.save(task)
        return {"critic_result": ensemble_result.__dict__}

    async def _pre_validate_patch(self, task: TaskState, state: WorkflowState) -> ToolResult:
        """Pre-validate a patch in a sandboxed environment before applying to the main codebase."""
        patch_text = state.get("patch_text", "")
        file_writes = state.get("file_writes", [])

        # Skip pre-validation for minor changes
        if not patch_text.strip() and not file_writes:
            return ToolResult(success=True, command="pre_validation", stdout="No changes to validate")

        issues: list[StructuredError] = []

        # Check for dangerous patterns in the patch
        dangerous_patterns = [
            (r'rm\s+-rf\s+/', "Destructive: 'rm -rf /' command detected"),
            (r"shutil\.rmtree\(", "Destructive: rmtree detected"),
            (r'os\.system\(["\'].*rm', "Destructive: os.system with rm command"),
            (r'subprocess\.call\(.*rm', "Destructive: subprocess call with rm"),
            (r'open\(["\"]w["\'].*\)', "File write: ensure proper encoding and error handling"),
            (r'eval\s*\(', "Security: eval() can execute arbitrary code"),
            (r'exec\s*\(', "Security: exec() can execute arbitrary code"),
            (r'__import__\s*\(["\']', "Dynamic import: potential security risk"),
        ]

        for pattern, description in dangerous_patterns:
            if re.search(pattern, patch_text, re.IGNORECASE):
                issues.append(
                    StructuredError(
                        error_type="SecurityRisk",
                        message=description,
                        file="patch",
                    )
                )

        # Check for syntax errors in patch
        if file_writes:
            for file_write in file_writes:
                try:
                    ast.parse(file_write.get("content", ""))
                except SyntaxError as e:
                    issues.append(
                        StructuredError(
                            error_type="SyntaxError",
                            message=f"Syntax error in {file_write.get('path')}: {e.msg}",
                            file=file_write.get("path"),
                            line=getattr(e, "lineno", None),
                        )
                    )

        if issues:
            return ToolResult(
                success=False,
                command="pre_validation",
                stderr="Pre-validation blocked dangerous or malformed patch",
                structured_errors=issues,
            )

        return ToolResult(success=True, command="pre_validation", stdout="Pre-validation passed")

    async def _apply_patch(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("apply_patch")
        self.store.save(task)

        # Pre-validate the patch before applying
        prevalidation_result = await self._pre_validate_patch(task, state)
        if not prevalidation_result.success:
            self._append_command(task, "pre_validation", prevalidation_result)
            self._append_step(task, "apply_patch", "PreValidator", "failed", prevalidation_result.stderr)
            task.error_message = f"Pre-validation failed: {prevalidation_result.stderr}"
            task.mark_terminal(TaskLifecycleStatus.failed, summary=task.error_message)
            self.store.save(task)
            return {"apply_result": ToolResult(
                success=False,
                command="pre_validation",
                stderr=prevalidation_result.stderr,
                structured_errors=prevalidation_result.structured_errors,
            ).model_dump(mode="json")}

        patch_text = state.get("patch_text", "")
        file_writes = state.get("file_writes", [])
        if file_writes:
            result = await self._apply_file_writes(task, file_writes)
        elif patch_text.strip():
            result = await self.git_runner.apply_patch(
                patch=patch_text,
                repo_path=task.execution_repo_path(),
                timeout_seconds=task.constraints.command_timeout_seconds,
            )
        else:
            result = ToolResult(success=True, stdout="No change proposal for this attempt.", command="apply change")
        if result.success:
            effective_diff_result = await self.git_runner.diff_repo(
                repo_path=task.execution_repo_path(),
                base_path=task.repo_path if task.workspace_strategy == "copy" else None,
            )
            effective_diff = effective_diff_result.stdout if effective_diff_result.success else ""
            if effective_diff.strip():
                result.data["effective_diff"] = effective_diff
            elif self._requires_code_changes(task.goal):
                result = ToolResult(
                    success=False,
                    command=result.command or "apply patch",
                    stdout=result.stdout,
                    stderr="Patch application produced no effective repository changes.",
                    exit_code=result.exit_code,
                    timed_out=result.timed_out,
                    started_at=result.started_at,
                    completed_at=result.completed_at,
                    data={**result.data, "no_effective_change": True},
                    structured_errors=result.structured_errors,
                )
        self._append_command(task, "apply_patch", result)
        status = "completed" if result.success else "failed"
        details = result.stderr.strip() or result.stdout.strip() or state.get("patch_rationale", "")
        if not result.success:
            task.error_message = details
        self._append_step(task, "apply_patch", "Executor", status, details)
        self._advance_subtask(task)
        self.store.save(task)
        return {"apply_result": result.model_dump(mode="json")}

    async def _run_validation(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("run_validation")
        self.store.save(task)

        validation_summary = await self.agent_controller.run_validation(task)
        for check in validation_summary.checks:
            self._append_command_from_validation(task, check)
            if check.name == "tests":
                self._append_test_result(task, check)
                self._add_artifact(
                    task,
                    self.artifact_store.write_text(
                        task.task_id,
                        f"attempt_{task.attempt}_test_log.txt",
                        "\n".join(filter(None, [check.stdout, check.stderr])),
                    ),
                )

        self._add_artifact(
            task,
            self.artifact_store.write_json(
                task.task_id,
                f"attempt_{task.attempt}_validation.json",
                validation_summary.model_dump(mode="json"),
            ),
        )
        step_status = "completed" if validation_summary.valid else "failed"
        details = "Validation passed." if validation_summary.valid else ", ".join(validation_summary.failed_checks)
        self._append_step(task, "run_validation", "Tester", step_status, details)
        self._advance_subtask(task)
        self.store.save(task)
        return {"validation_result": validation_summary.model_dump(mode="json")}

    async def _evaluate_result(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("evaluate_result")
        self.store.save(task)

        apply_result = ToolResult.model_validate(state.get("apply_result", {"success": True}))
        validation_summary = ValidationSummary.model_validate(
            state.get("validation_result", {"valid": True, "failed_checks": [], "checks": []})
        )

        critic_result, evaluation_summary = self.agent_controller.evaluate(
            task,
            validation_summary,
            patch=state.get("patch_text", ""),
            file_writes=state.get("file_writes", []),
        )
        issues = list(evaluation_summary.issues)
        focus_review = await self.model_router.review_focus_fix(
            task=task,
            repo_snapshot=task.repo_context,
            validation_summary=validation_summary,
            instruction=self.prompt_loader.load("repair_reviewer.txt"),
        )
        task.focus_review_history.append(focus_review)
        if task.focus_issue_inventory is not None:
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    f"attempt_{task.attempt}_focus_review.json",
                    focus_review.model_dump(mode="json"),
                ),
            )
        if not focus_review.passed:
            issues.extend(f"Focused review unresolved: {item}" for item in focus_review.unresolved_issues)
        failure_context = None
        success = apply_result.success and evaluation_summary.passed and focus_review.passed

        if not apply_result.success:
            issues.insert(0, apply_result.stderr or apply_result.stdout or "Patch application failed.")

        validation_feedback = self._format_validation_feedback(validation_summary)
        if not success:
            retry_focus_files = self._collect_failure_focus_files(task, validation_summary)
            retry_test_commands = self._collect_retry_test_commands(task, validation_summary)
            if retry_focus_files:
                task.retry_focus_files = retry_focus_files
                task.repo_context["failure_focus_files"] = retry_focus_files
                self._expand_repo_snapshot(task=task, repo_snapshot=task.repo_context, requested_paths=retry_focus_files)
            if retry_test_commands:
                task.repo_context["retry_test_commands"] = retry_test_commands
            test_check = next(
                (check for check in validation_summary.checks if check.name.startswith("tests") and not check.success),
                next((check for check in validation_summary.checks if check.name.startswith("tests")), None),
            )
            if not focus_review.passed and focus_review.unresolved_issues:
                failure_context = "Focused file review found remaining issues:\n- " + "\n- ".join(focus_review.unresolved_issues)
                if validation_feedback:
                    failure_context = f"{failure_context}\n\nExact validation diagnostics:\n{validation_feedback}"
            elif test_check is not None and not test_check.success:
                analyzed_failure = await self.model_router.analyze_failure(
                    task=task,
                    test_stdout=test_check.stdout,
                    test_stderr=test_check.stderr,
                    exit_code=test_check.exit_code,
                    instruction=self.prompt_loader.load("critic.txt"),
                )
                failure_context = analyzed_failure
                if validation_feedback:
                    failure_context = f"{failure_context}\n\nExact validation diagnostics:\n{validation_feedback}"
            elif validation_feedback:
                failure_context = (
                    "Validation failed. Fix the exact diagnostics below before changing anything else:\n"
                    f"{validation_feedback}"
                )
            else:
                failure_context = issues[0] if issues else "Validation failed."
            if critic_result.repair_guidance:
                guidance_lines = self._dedupe_text_lines(critic_result.repair_guidance)[:8]
                if guidance_lines:
                    failure_context = (
                        f"{failure_context}\n\nRepair guidance:\n- " + "\n- ".join(guidance_lines)
                    )
            patterns = await self.agent_controller.memory_service.retrieve_patterns(
                failure_context,
                overrides=task.constraints.model_overrides,
            )
            if patterns:
                failure_context = f"{failure_context}\nRelated patterns: {' | '.join(patterns)}"
        else:
            task.retry_focus_files = []
            task.repo_context.pop("failure_focus_files", None)
            task.repo_context.pop("retry_test_commands", None)

        should_retry = (
            not success
            and task.attempt < task.constraints.max_attempts
            and not self.model_router.is_mock_step("coding", task.constraints.model_overrides)
        )
        evaluation = {
            "success": success,
            "should_retry": should_retry,
            "reason": failure_context or "Validation passed.",
            "issues": issues,
            "validation_feedback": validation_feedback.splitlines() if validation_feedback else [],
            "repair_guidance": self._dedupe_text_lines(critic_result.repair_guidance),
        }
        task.error_message = None if success else evaluation["reason"]
        task.reflection_issues = critic_result.issues
        self._add_artifact(
            task,
            self.artifact_store.write_json(
                task.task_id,
                f"attempt_{task.attempt}_evaluation.json",
                evaluation,
            ),
        )
        details = evaluation["reason"] if not success else "Evaluation passed."
        self._append_step(task, "evaluate_result", "Critic", "completed", details)
        self._advance_subtask(task)
        self.store.save(task)
        return {
            "evaluation": evaluation,
            "failure_context": failure_context,
            "critic_result": critic_result.model_dump(mode="json"),
        }

    async def _summarize(self, state: WorkflowState) -> dict[str, Any]:
        task = self._require_task(state["task_id"])
        task.mark_running("summarize")
        self.store.save(task)

        evaluation = state.get("evaluation", {})
        success = bool(evaluation.get("success", True))
        final_diff = await self.agent_controller.finalize_task(task, success=success)
        if success and self._requires_code_changes(task.goal) and not final_diff.strip():
            success = False
            task.error_message = "No effective code changes were committed for a fix-oriented task."
        if final_diff.strip():
            self._add_artifact(task, self.artifact_store.write_text(task.task_id, "final.diff", final_diff))

        if task.status != TaskLifecycleStatus.cancelled:
            task.status = TaskLifecycleStatus.completed if success else TaskLifecycleStatus.failed
        if task.status == TaskLifecycleStatus.failed and not task.error_message:
            task.error_message = str(evaluation.get("reason", "Task failed."))

        summary = await self.model_router.summarize(
            task=task,
            final_diff=final_diff,
            instruction=self.prompt_loader.load("summarizer.txt"),
        )
        aggregate = self.agent_controller.supervisor_agent.aggregate_results(
            plan_steps=task.plan_steps,
            failed_checks=task.validation_history[-1].failed_checks if task.validation_history else [],
            issues=list(evaluation.get("issues", [])),
        )
        task.mark_terminal(task.status, summary=summary, error_message=task.error_message)
        learned_example_ids: list[str] = []
        if success and final_diff.strip():
            learned_example_ids = self.agent_controller.promote_successful_fix(task, final_diff=final_diff)
        self._add_artifact(task, self.artifact_store.write_text(task.task_id, "summary.txt", summary))
        self._add_artifact(task, self.artifact_store.write_json(task.task_id, "aggregate_result.json", aggregate))
        if learned_example_ids:
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    "rag_learning_examples.json",
                    {"count": len(learned_example_ids), "example_ids": learned_example_ids},
                ),
            )
        if task.rag_retrieval_events:
            self._add_artifact(
                task,
                self.artifact_store.write_json(
                    task.task_id,
                    "rag_retrieval_events.json",
                    task.rag_retrieval_events,
                ),
            )
        self._append_step(task, "summarize", "Doc Agent", "completed", "Summary generated.")
        self.store.save(task)
        return {"summary": summary, "final_diff": final_diff}

    def _route_after_critic(self, state: WorkflowState) -> str:
        critic_result = CriticResult.model_validate(state.get("critic_result", {"result": "pass"}))
        task = self._require_task(state["task_id"])
        if (
            critic_result.result == "fail"
            and task.attempt < task.constraints.max_attempts
            and not self.model_router.is_mock_step("coding", task.constraints.model_overrides)
        ):
            return "retry"
        return "apply"

    def _route_after_evaluation(self, state: WorkflowState) -> str:
        evaluation = state.get("evaluation", {})
        if evaluation.get("should_retry"):
            return "retry"
        return "summarize"

    async def _apply_file_writes(self, task: TaskState, file_writes: list[dict[str, str]]) -> ToolResult:
        written_paths: list[str] = []
        repo_root = task.execution_repo_path().resolve()
        for file_write in file_writes:
            relative_path = self._normalize_relative_path(str(file_write.get("path", "")))
            if relative_path is None:
                return ToolResult(
                    success=False,
                    stderr=f"Invalid file write path: {file_write.get('path', '')}",
                    command="rewrite files",
                )
            target_path = (repo_root / relative_path).resolve()
            try:
                target_path.relative_to(repo_root)
            except ValueError:
                return ToolResult(
                    success=False,
                    stderr=f"File write escapes repository root: {relative_path.as_posix()}",
                    command="rewrite files",
                )
            if target_path.exists() and target_path.is_dir():
                return ToolResult(
                    success=False,
                    stderr=f"File write targets a directory, not a file: {relative_path.as_posix()}",
                    command="rewrite files",
                )

            result = self.file_tools.write_file(target_path, str(file_write.get("content", "")))
            if not result.success:
                result.command = "rewrite files"
                return result
            written_paths.append(relative_path.as_posix())

        return ToolResult(
            success=True,
            stdout=f"Rewrote {len(written_paths)} file(s).",
            command="rewrite files",
            data={"written_paths": written_paths},
        )

    def _expand_repo_snapshot(
        self,
        *,
        task: TaskState,
        repo_snapshot: dict[str, Any],
        requested_paths: list[str],
    ) -> list[str]:
        expanded: list[str] = []
        repo_root = task.execution_repo_path().resolve()
        previews = list(repo_snapshot.get("file_previews", []))
        preview_index = {
            str(item.get("path", "")).strip(): index
            for index, item in enumerate(previews)
            if isinstance(item, dict)
        }
        mentioned_files = [str(path) for path in repo_snapshot.get("mentioned_files", [])]
        important_files = [str(path) for path in repo_snapshot.get("important_files", [])]
        key_modules = [str(path) for path in repo_snapshot.get("key_modules", [])]

        for raw_path in requested_paths:
            relative_path = self._normalize_relative_path(raw_path)
            if relative_path is None:
                continue
            target_path = (repo_root / relative_path).resolve()
            try:
                target_path.relative_to(repo_root)
            except ValueError:
                continue
            read_result = self.file_tools.read_file(target_path)
            if not read_result.success:
                continue
            content = read_result.stdout
            relative_str = relative_path.as_posix()
            preview_payload = {
                "path": relative_str,
                "content": content[: self.settings.repo_focus_file_bytes],
                "truncated": len(content) > self.settings.repo_focus_file_bytes,
            }
            if preview_payload["truncated"]:
                preview_payload["content"] = f"{preview_payload['content']}\n...<truncated>"
            if relative_str in preview_index:
                previews[preview_index[relative_str]] = preview_payload
            else:
                previews.append(preview_payload)
            for collection in (mentioned_files, important_files, key_modules):
                if relative_str not in collection:
                    collection.append(relative_str)
            expanded.append(relative_str)

        repo_snapshot["file_previews"] = previews
        repo_snapshot["mentioned_files"] = mentioned_files
        repo_snapshot["important_files"] = important_files
        repo_snapshot["key_modules"] = key_modules
        return expanded

    def _get_attempt_instruction(self, attempt: int, max_attempts: int) -> str:
        """Get progressive retry strategy based on attempt number."""
        base_instruction = self.prompt_loader.load("coder.txt")

        if attempt == 1:
            return base_instruction
        if attempt == 2:
            return f"""{base_instruction}

CRITICAL - You are on attempt {attempt} of {max_attempts}.

Previous validation details are in the failure_context payload. Treat those exact diagnostics as the highest priority.

CHANGE YOUR APPROACH:
- Focus on MINIMAL changes only
- Clear concrete linter and validator errors before additional refactors
- Prefer the smallest rewrite that fully removes the reported failure
- Do not ignore exact file/line diagnostics

Think carefully about what went wrong and try a completely different approach.
"""
        remaining_attempts = max_attempts - attempt
        final_note = (
            "This is the final attempt. Submit the cleanest complete fix you can."
            if remaining_attempts <= 0
            else f"There are {remaining_attempts} retry slot(s) remaining after this attempt."
        )
        return f"""{base_instruction}

CRITICAL - You are on attempt {attempt} of {max_attempts}.
{final_note}

Use a more decisive recovery strategy:
- Prefer mode="rewrite_files" for a single focused Python file you fully possess
- Resolve every exact validation diagnostic from failure_context before anything else
- Eliminate lint leftovers and false partial fixes
- Keep the result readable and validation-clean
"""

    @staticmethod
    def _dedupe_text_lines(lines: list[str]) -> list[str]:
        seen: set[str] = set()
        unique: list[str] = []
        for line in lines:
            cleaned = str(line).strip()
            if not cleaned or cleaned in seen:
                continue
            seen.add(cleaned)
            unique.append(cleaned)
        return unique

    def _format_validation_feedback(self, validation_summary: ValidationSummary, *, limit: int | None = None) -> str:
        limit = limit or self.settings.validation_feedback_max_items
        formatted: list[str] = []
        for check in validation_summary.checks:
            if check.success:
                continue
            if check.structured_errors:
                for structured_error in check.structured_errors:
                    location = structured_error.file or "unknown file"
                    if structured_error.line:
                        location = f"{location}:{structured_error.line}"
                    detail = structured_error.message or structured_error.error_type
                    formatted.append(f"- {check.name}: {structured_error.error_type} at {location}: {detail}")
                    if len(formatted) >= limit:
                        return "\n".join(formatted)
            else:
                detail = check.stderr or check.stdout or "validation check failed"
                formatted.append(f"- {check.name}: {detail}")
                if len(formatted) >= limit:
                    return "\n".join(formatted)
        return "\n".join(formatted)

    @staticmethod
    def _normalize_relative_path(raw_path: str) -> Path | None:
        cleaned = raw_path.strip().replace("\\", "/").lstrip("./")
        if not cleaned:
            return None
        relative_path = Path(cleaned)
        if relative_path.is_absolute() or ".." in relative_path.parts:
            return None
        return relative_path

    def _require_task(self, task_id: str) -> TaskState:
        task = self.store.get(task_id)
        if task is None:
            raise LookupError(f"Unknown task: {task_id}")
        return task

    def _append_step(self, task: TaskState, step: str, role: str, status: str, details: str | None) -> None:
        task.step_history.append(StepRecord(step=step, role=role, status=status, details=details))

    def _append_command(self, task: TaskState, tool_name: str, result: ToolResult) -> None:
        task.command_history.append(
            CommandRecord(
                tool=tool_name,
                command=result.command or "",
                success=result.success,
                stdout=result.stdout,
                stderr=result.stderr,
                exit_code=result.exit_code,
                structured_errors=result.structured_errors,
                started_at=result.started_at,
                completed_at=result.completed_at,
            )
        )

    def _append_command_from_validation(self, task: TaskState, check) -> None:
        task.command_history.append(
            CommandRecord(
                tool=f"validation:{check.name}",
                command=check.command or "",
                success=check.success,
                stdout=check.stdout,
                stderr=check.stderr,
                exit_code=check.exit_code,
                structured_errors=check.structured_errors,
            )
        )

    def _append_test_result(self, task: TaskState, check) -> None:
        task.test_results.append(
            TestResult(
                command=check.command,
                success=check.success,
                stdout=check.stdout,
                stderr=check.stderr,
                exit_code=check.exit_code,
                structured_errors=check.structured_errors,
            )
        )

    def _advance_subtask(self, task: TaskState) -> None:
        for subtask in task.subtasks:
            if subtask.status == "pending":
                subtask.status = "completed"
                return

    def _add_artifact(self, task: TaskState, artifact: ArtifactRecord) -> None:
        if any(existing.path == artifact.path and existing.type == artifact.type for existing in task.artifacts):
            return
        task.artifacts.append(artifact)

    @staticmethod
    def _requires_code_changes(goal: str) -> bool:
        lowered = goal.lower()
        return any(keyword in lowered for keyword in ("fix", "repair", "resolve", "patch", "correct", "debug"))

    async def _retrieve_rag_examples(self, task: TaskState, repo_snapshot: dict) -> list[dict[str, Any]]:
        """Retrieve similar bug-fix examples from the RAG knowledge base."""
        if not self.rag_knowledge_base or not self.embedding_service:
            return []

        # Build query from task goal and code context
        mentioned_files = repo_snapshot.get("mentioned_files", [])
        if not mentioned_files:
            return []

        examples: list[dict[str, Any]] = []
        seen_example_ids: set[str] = set()

        # For each mentioned file, retrieve similar examples
        for file_path in mentioned_files[: self.settings.rag_task_file_limit]:
            try:
                # Read the file content
                full_path = task.execution_repo_path() / file_path
                if not full_path.exists():
                    continue

                content = full_path.read_text(encoding="utf-8")

                for unit in self.rag_knowledge_base.extract_query_units(content, file_path=str(file_path)):
                    query_code = str(unit.get("code", "")).strip()
                    if not query_code:
                        continue
                    code_embedding = self.embedding_service.embed_code(query_code, context=task.goal)
                    query = self.rag_knowledge_base.build_retrieval_query(
                        query_code=query_code,
                        bug_description=task.goal,
                        file_path=str(file_path),
                        function_name=unit.get("function_name"),
                        max_results=self.settings.rag_task_max_results_per_unit,
                    )
                    result = self.rag_knowledge_base.retrieve_similar(
                        query=query,
                        query_embedding=code_embedding,
                    )
                    if result.examples:
                        event_id = self.rag_knowledge_base.record_retrieval_event(
                            task_id=task.task_id,
                            stage="coding",
                            query=query,
                            examples=result.examples,
                            file_path=str(file_path),
                            function_name=unit.get("function_name"),
                        )
                        task.rag_retrieval_events.append(
                            {
                                "event_id": event_id,
                                "stage": "coding",
                                "file_path": str(file_path),
                                "function_name": unit.get("function_name"),
                                "example_ids": [item.example.id for item in result.examples],
                            }
                        )

                    for retrieved in result.examples:
                        if retrieved.relevance_score <= self.settings.rag_task_min_relevance_score or retrieved.example.id in seen_example_ids:
                            continue
                        seen_example_ids.add(retrieved.example.id)
                        examples.append(
                            {
                                "summary": retrieved.example.bug_summary,
                                "category": retrieved.example.bug_category,
                                "repository": retrieved.example.repository,
                                "example_kind": retrieved.example.example_kind,
                                "file_path": retrieved.example.file_path,
                                "function_name": retrieved.example.function_name,
                                "buggy_code_snippet": retrieved.example.buggy_code[:500],
                                "fixed_code_snippet": retrieved.example.fixed_code[:500],
                                "diff_snippet": retrieved.example.diff[:500],
                                "relevance_score": retrieved.relevance_score,
                                "suggested_patterns": retrieved.example.related_patterns,
                            }
                        )

            except (OSError, UnicodeDecodeError):
                continue

        examples.sort(key=lambda item: item["relevance_score"], reverse=True)
        return examples[: self.settings.rag_task_total_examples_limit]

    def _collect_failure_focus_files(self, task: TaskState, validation_summary: ValidationSummary) -> list[str]:
        focus_files: list[str] = []
        seen: set[str] = set()
        for check in validation_summary.checks:
            for structured_error in check.structured_errors:
                for raw_path in [structured_error.file, *structured_error.metadata.get("traceback_files", [])]:
                    relative_path = self._normalize_repo_relative_path(task, raw_path)
                    if relative_path is None or relative_path in seen:
                        continue
                    seen.add(relative_path)
                    focus_files.append(relative_path)
        return focus_files[:6]

    def _collect_retry_test_commands(self, task: TaskState, validation_summary: ValidationSummary) -> list[str]:
        commands: list[str] = []
        seen: set[str] = set()
        for check in validation_summary.checks:
            if not check.name.startswith("tests"):
                continue
            for structured_error in check.structured_errors:
                relative_path = self._normalize_repo_relative_path(task, structured_error.file)
                if relative_path is None or not relative_path.endswith(".py"):
                    continue
                command = ""
                if structured_error.test_name:
                    command = f"pytest -q {relative_path}::{structured_error.test_name}"
                elif relative_path.startswith("tests/") or Path(relative_path).name.startswith("test_"):
                    command = f"pytest -q {relative_path}"
                if command and command not in seen:
                    seen.add(command)
                    commands.append(command)
        return commands[:4]

    def _normalize_repo_relative_path(self, task: TaskState, raw_path: str | None) -> str | None:
        if not raw_path:
            return None
        normalized = str(raw_path).strip().replace("\\", "/")
        if not normalized:
            return None
        candidate = Path(normalized)
        roots = [task.execution_repo_path().resolve(), task.repo_path.resolve()]
        if candidate.is_absolute():
            try:
                resolved = candidate.resolve()
            except OSError:
                return None
            for root in roots:
                try:
                    return resolved.relative_to(root).as_posix()
                except ValueError:
                    continue
            return None

        cleaned = normalized.lstrip("./")
        relative = Path(cleaned)
        if relative.is_absolute() or ".." in relative.parts:
            return None
        return relative.as_posix()
