from __future__ import annotations

from pathlib import Path

from agents.coder import Coder
from agents.judge import Judge
from agents.reviewer import Reviewer
from control_plane.agents.models import AgentRole, AgentVersion
from control_plane.events.bus import EventBus
from control_plane.events.models import EventType
from control_plane.projects.models import CommitRecord, Task, TaskAttempt, TaskStatus, Worktree
from control_plane.verification.models import VerificationResult
from graph.scheduler import TaskScheduler
from knowledge.context_builder import WorkerContextBuilder
from model_router.router import ModelRouter
from tools.capabilities import Capability
from tools.runtime import WorktreeTools
from tools.test_runner import DeterministicTestRunner
from workspace.worktrees import WorktreeManager


class AutonomousTaskLoop:
    def __init__(
        self,
        router: ModelRouter,
        *,
        scheduler: TaskScheduler | None = None,
        bus: EventBus | None = None,
    ) -> None:
        self.router = router
        self.bus = bus or EventBus()
        self.scheduler = scheduler or TaskScheduler(self.bus)
        self.context_builder = WorkerContextBuilder()
        self.worktrees = WorktreeManager()
        self.coder = Coder(router)
        self.reviewer = Reviewer()
        self.judge = Judge()
        self.tests = DeterministicTestRunner()

    def run_once(self, *, test_command: list[str] | None = None) -> Task | None:
        task = self.scheduler.claim_next_ready_task()
        if task is None:
            return None
        self._execute_task(task, test_command or ["python", "-m", "pytest"])
        return task

    def _execute_task(self, task: Task, test_command: list[str]) -> None:
        coder_version = self._champion(AgentRole.CODER)
        reviewer_version = self._champion(AgentRole.REVIEWER)
        judge_version = self._champion(AgentRole.PROJECT_JUDGE)
        worktree = self._get_or_create_worktree(task)
        tools = WorktreeTools(
            Path(worktree.worktree_path),
            {
                Capability.READ_REPOSITORY,
                Capability.WRITE_WORKTREE,
                Capability.RUN_TESTS,
                Capability.COMMIT_CHANGES,
            },
        )

        while task.retry_count <= task.max_retries:
            attempt = TaskAttempt.objects.create(
                task=task,
                attempt_number=task.retry_count + 1,
                coder=coder_version,
                status="RUNNING",
            )
            context = self.context_builder.build_for_task(task, Path(worktree.worktree_path))
            context["previous_attempts"] = list(
                task.attempts.exclude(id=attempt.id).order_by("attempt_number").values(
                    "attempt_number", "status", "coder_result", "review_findings", "judge_findings"
                )
            )
            attempt.context_snapshot = self._scrub_context(context)
            attempt.save(update_fields=["context_snapshot", "updated_at"])

            coder_result = self.coder.execute(context, tools, project=task.project, agent_version=coder_version)
            attempt.coder_result = {
                "status": coder_result.status,
                "summary": coder_result.summary,
                "changed_files": coder_result.changed_files,
                "metadata": coder_result.metadata,
            }
            attempt.save(update_fields=["coder_result", "updated_at"])
            if coder_result.status != "COMPLETE":
                if self._retry_or_fail(task, attempt, "coder_failed", [coder_result.summary]):
                    continue
                return

            test_run = self.tests.run(task.project, task, Path(worktree.worktree_path), test_command)
            if test_run.status != "PASS":
                self.bus.publish(EventType.TEST_FAILED, project=task.project, task=task, payload={"test_run_id": str(test_run.id)})

            tools.git(["add", "-N", "."])
            diff = tools.diff()
            review = self.reviewer.review(task, reviewer_version, diff, test_run.status)
            attempt.review_findings = review.findings
            attempt.save(update_fields=["review_findings", "updated_at"])
            if review.status != "PASS":
                self.bus.publish(
                    EventType.REVIEW_FAILED,
                    project=task.project,
                    task=task,
                    payload={"review_id": str(review.id), "findings": review.findings},
                )
                if self._retry_or_fail(task, attempt, "review_failed", review.findings):
                    continue
                return

            verification = self.judge.judge(task.project, task, judge_version, diff, test_run.status)
            attempt.judge_findings = verification.evidence
            attempt.save(update_fields=["judge_findings", "updated_at"])
            if verification.result != VerificationResult.PASS:
                if self._retry_or_fail(task, attempt, "judge_failed", verification.evidence):
                    continue
                return

            sha = tools.commit_all(f"Artifex task: {task.goal[:80]}")
            commit = CommitRecord.objects.create(
                project=task.project,
                task=task,
                worktree=worktree,
                coder=coder_version,
                reviewer=reviewer_version,
                judge=judge_version,
                test_run=test_run,
                review=review,
                verification=verification,
                sha=sha,
                branch_name=worktree.branch_name,
                message=f"Artifex task: {task.goal[:80]}",
            )
            attempt.status = "COMPLETE"
            attempt.save(update_fields=["status", "updated_at"])
            task.status = TaskStatus.COMPLETE
            task.save(update_fields=["status", "updated_at"])
            self.bus.publish(EventType.COMMIT_CREATED, project=task.project, task=task, payload={"commit_id": str(commit.id), "sha": sha})
            self.bus.publish(EventType.TASK_COMPLETED, project=task.project, task=task, payload={"task_id": str(task.id)})
            self.worktrees.validate_clean_worktree(worktree)
            self.worktrees.cleanup(worktree)
            return

    def _get_or_create_worktree(self, task: Task) -> Worktree:
        try:
            return task.worktree
        except Worktree.DoesNotExist:
            if not task.project.repository_path:
                raise RuntimeError("Task project has no repository_path")
            return self.worktrees.create_for_task(task, Path(task.project.repository_path))

    def _champion(self, role: AgentRole) -> AgentVersion:
        return AgentVersion.objects.select_related("agent").get(agent__role=role, promotion_status="CHAMPION")

    def _retry_or_fail(self, task: Task, attempt: TaskAttempt, reason: str, findings: object) -> bool:
        attempt.status = "REWORK_REQUIRED" if task.retry_count < task.max_retries else "FAILED"
        attempt.save(update_fields=["status", "updated_at"])
        task.retry_count += 1
        if task.retry_count <= task.max_retries:
            task.status = TaskStatus.RUNNING
            task.save(update_fields=["retry_count", "status", "updated_at"])
            self.bus.publish(EventType.TASK_FAILED, project=task.project, task=task, payload={"reason": reason, "will_retry": True, "findings": findings})
            return True
        task.status = TaskStatus.FAILED
        task.save(update_fields=["retry_count", "status", "updated_at"])
        self.bus.publish(EventType.TASK_FAILED, project=task.project, task=task, payload={"reason": reason, "will_retry": False, "findings": findings})
        return False

    def _scrub_context(self, context: dict[str, object]) -> dict[str, object]:
        scrubbed = dict(context)
        scrubbed.pop("secrets", None)
        return scrubbed