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.bootstrap import champion_task_execution_graph_v1
from graph.langgraph_runtime import LangGraphRuntime
from graph.models import GraphRun
from graph.scheduler import TaskScheduler
from graph.task_nodes import TaskExecutionServices, task_execution_registry
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:
        graph_version = champion_task_execution_graph_v1()
        graph_run = GraphRun.objects.create(
            execution_graph_version=graph_version,
            project=task.project,
            milestone=task.milestone,
            feature=task.feature,
            task=task,
            current_node=graph_version.graph_spec["entry"],
        )
        services = TaskExecutionServices(self.router, bus=self.bus, test_command=test_command)
        LangGraphRuntime(task_execution_registry(services), bus=self.bus).run_until_terminal_or_paused(graph_run)
        return

        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.INVESTIGATE_WORKTREE,
                Capability.WRITE_WORKTREE,
                Capability.RUN_TESTS,
                Capability.COMMIT_CHANGES,
            },
        )

        effective_max_retries = min(task.max_retries, 2)
        while task.retry_count <= effective_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], effective_max_retries):
                    continue
                return

            test_run = self.tests.run(task.project, task, Path(worktree.worktree_path), test_command)
            if test_run.status != "PASS":
                self._attach_test_failure_evidence(attempt, test_run)
                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, effective_max_retries):
                    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, effective_max_retries):
                    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"])
            telemetry = self._task_telemetry(task)
            self.bus.publish(EventType.COMMIT_CREATED, project=task.project, task=task, payload={"commit_id": str(commit.id), "sha": sha, "telemetry": telemetry})
            self.bus.publish(EventType.TASK_COMPLETED, project=task.project, task=task, payload={"task_id": str(task.id), "telemetry": telemetry})
            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, effective_max_retries: int) -> bool:
        attempt.status = "REWORK_REQUIRED" if task.retry_count < effective_max_retries else "FAILED"
        attempt.save(update_fields=["status", "updated_at"])
        task.retry_count += 1
        classification = self._classify_failure(reason, findings)
        if task.retry_count <= effective_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, "classification": classification, "findings": findings})
            return True
        task.status = TaskStatus.FAILED
        task.save(update_fields=["retry_count", "status", "updated_at"])
        self.bus.publish("TASK_RETRY_EXHAUSTED", project=task.project, task=task, payload={"reason": reason, "classification": classification, "findings": findings})
        self.bus.publish(EventType.TASK_FAILED, project=task.project, task=task, payload={"reason": reason, "will_retry": False, "classification": classification, "findings": findings})
        return False

    def _classify_failure(self, reason: str, findings: object) -> str:
        text = f"{reason} {findings}".lower()
        if "unsupported operation" in text or "missing capability" in text:
            return "missing_capability"
        if "context" in text or "migration" in text:
            return "context_problem"
        if "timeout" in text or "provider" in text:
            return "environment_problem"
        if "malformed json" in text or "model" in text:
            return "model_problem"
        if "ambiguous" in text:
            return "intent_ambiguity"
        if reason == "review_failed" or reason == "judge_failed":
            return "replan"
        return "split_task"

    def _task_telemetry(self, task: Task) -> dict[str, object]:
        totals: dict[str, float] = {
            "patch_mismatch_count": 0,
            "patch_attempts": 0,
            "patch_successes": 0,
            "write_file_operations": 0,
            "write_file_fallbacks": 0,
            "mutation_operations_per_accepted_task": 0,
            "model_requests_per_task": 0,
        }
        for attempt in task.attempts.all():
            metadata = attempt.coder_result.get("metadata", {}) if isinstance(attempt.coder_result, dict) else {}
            telemetry = metadata.get("telemetry", {}) if isinstance(metadata, dict) else {}
            if not isinstance(telemetry, dict):
                continue
            totals["patch_mismatch_count"] += float(telemetry.get("patch_mismatches", 0))
            totals["patch_attempts"] += float(telemetry.get("patch_attempts", 0))
            totals["patch_successes"] += float(telemetry.get("patch_successes", 0))
            totals["write_file_operations"] += float(telemetry.get("write_file_operations", 0))
            totals["write_file_fallbacks"] += float(telemetry.get("write_file_fallbacks", 0))
            totals["mutation_operations_per_accepted_task"] += float(telemetry.get("mutation_operations", 0))
            totals["model_requests_per_task"] += float(telemetry.get("model_requests", 0))
        patch_attempts = totals["patch_attempts"]
        write_file_operations = totals["write_file_operations"]
        return {
            **{key: int(value) for key, value in totals.items()},
            "patch_success_rate": 0 if patch_attempts == 0 else totals["patch_successes"] / patch_attempts,
            "write_file_fallback_rate": 0 if write_file_operations == 0 else totals["write_file_fallbacks"] / write_file_operations,
        }

    def _attach_test_failure_evidence(self, attempt: TaskAttempt, test_run) -> None:
        content = test_run.output_artifact.content if test_run.output_artifact else {}
        stdout = str(content.get("stdout", "")) if isinstance(content, dict) else ""
        stderr = str(content.get("stderr", "")) if isinstance(content, dict) else ""
        evidence = {
            "test_run_id": str(test_run.id),
            "status": test_run.status,
            "stdout_excerpt": stdout[-12000:],
            "stderr_excerpt": stderr[-4000:],
        }
        coder_result = dict(attempt.coder_result or {})
        metadata = dict(coder_result.get("metadata", {})) if isinstance(coder_result.get("metadata", {}), dict) else {}
        metadata["test_failure_evidence"] = evidence
        coder_result["metadata"] = metadata
        attempt.coder_result = coder_result
        attempt.save(update_fields=["coder_result", "updated_at"])

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