from __future__ import annotations

from pathlib import Path

from ai_orchestrator.agents.critic_agent import CriticAgent
from ai_orchestrator.agents.evaluator_agent import EvaluatorAgent
from ai_orchestrator.agents.planner_agent import PlannerAgent
from ai_orchestrator.schemas.task_state import TaskState
from ai_orchestrator.schemas.workflow_models import CriticResult, EvaluationSummary, ValidationSummary
from ai_orchestrator.services.memory_service import MemoryService
from ai_orchestrator.services.embedding_service import EmbeddingService
from ai_orchestrator.services.rag_knowledge_base import RAGKnowledgeBase
from ai_orchestrator.services.repo_analyzer import RepoAnalyzer
from ai_orchestrator.services.supervisor_agent import SupervisorAgent
from ai_orchestrator.services.validation_engine import ValidationEngine
from ai_orchestrator.services.workspace_manager import WorkspaceManager


class AgentController:
    def __init__(
        self,
        *,
        planner_agent: PlannerAgent,
        critic_agent: CriticAgent,
        evaluator_agent: EvaluatorAgent,
        supervisor_agent: SupervisorAgent,
        memory_service: MemoryService,
        repo_analyzer: RepoAnalyzer,
        validation_engine: ValidationEngine,
        workspace_manager: WorkspaceManager,
        rag_knowledge_base: RAGKnowledgeBase | None = None,
        embedding_service: EmbeddingService | None = None,
    ) -> None:
        self.planner_agent = planner_agent
        self.critic_agent = critic_agent
        self.evaluator_agent = evaluator_agent
        self.supervisor_agent = supervisor_agent
        self.memory_service = memory_service
        self.repo_analyzer = repo_analyzer
        self.validation_engine = validation_engine
        self.workspace_manager = workspace_manager
        self.rag_knowledge_base = rag_knowledge_base
        self.embedding_service = embedding_service

    async def create_workspace(self, task: TaskState) -> None:
        snapshot = await self.workspace_manager.create_snapshot(task.task_id, task.repo_path)
        task.workspace_path = Path(snapshot.workspace_path)
        task.workspace_strategy = snapshot.strategy

    async def analyze_repository(self, task: TaskState) -> dict:
        client_context = task.constraints.extra.get("client_context")
        repo_context = self.repo_analyzer.scan_repository(
            task.execution_repo_path(),
            goal=task.goal,
            client_context=client_context if isinstance(client_context, dict) else None,
        )
        repo_context["task_mode"] = task.task_mode
        task.repo_context = repo_context
        await self.memory_service.store_repo_knowledge(
            repo_context,
            overrides=task.constraints.model_overrides,
        )
        task.memory_hits = await self.memory_service.retrieve_similar_tasks(
            task.goal,
            task_mode=task.task_mode,
            focus_files=list(repo_context.get("mentioned_files", [])),
            overrides=task.constraints.model_overrides,
        )
        return repo_context

    async def plan_task(self, task: TaskState, instruction: str | None = None) -> list[str]:
        if task.task_mode == "focused_file_repair":
            focus_targets = task.repo_context.get("mentioned_files", []) or task.repo_context.get("important_files", [])
            target_label = ", ".join(focus_targets[:2]) if focus_targets else "the focused file"
            task.plan_steps = [
                f"Inspect {target_label} and enumerate concrete correctness issues.",
                f"Rewrite {target_label} to resolve every identified issue with the smallest safe change.",
                "Run syntax, runtime, and synthesized focused acceptance checks.",
                "Summarize the final diff and any unresolved issues.",
            ]
            task.subtasks = self.supervisor_agent.schedule_subtasks(
                self.supervisor_agent.create_subtasks(task.task_id, task.plan_steps)
            )
            return task.plan_steps
        if task.task_mode == "failing_test_repair":
            task.plan_steps = [
                "Inspect the failing test or validation path and isolate the breakage.",
                "Make the smallest code change that resolves the failure.",
                "Rerun validation and summarize the resulting diff.",
            ]
            task.subtasks = self.supervisor_agent.schedule_subtasks(
                self.supervisor_agent.create_subtasks(task.task_id, task.plan_steps)
            )
            return task.plan_steps

        memory_context = [self._format_memory_hit(hit) for hit in task.memory_hits]
        task.plan_steps = await self.planner_agent.create_plan(
            user_goal=task.goal,
            repo_context=task.repo_context,
            test_command=task.constraints.test_command,
            memory_context=memory_context,
            instruction=instruction,
            overrides=task.constraints.model_overrides,
        )
        task.subtasks = self.supervisor_agent.schedule_subtasks(
            self.supervisor_agent.create_subtasks(task.task_id, task.plan_steps)
        )
        return task.plan_steps

    def reflect_patch(
        self,
        task: TaskState,
        patch: str,
        *,
        file_writes: list[dict] | None = None,
        instruction: str | None = None,
    ) -> CriticResult:
        critic_result = self.critic_agent.review(
            goal=task.goal,
            patch=patch,
            file_writes=file_writes,
            repo_context=task.repo_context,
            validation_result=None,
        )
        task.reflection_issues = critic_result.issues
        return critic_result

    async def run_validation(self, task: TaskState) -> ValidationSummary:
        validation_summary = await self.validation_engine.validate(task, task.repo_context)
        task.validation_history.append(validation_summary)
        return validation_summary

    def evaluate(
        self,
        task: TaskState,
        validation_summary: ValidationSummary,
        *,
        patch: str,
        file_writes: list[dict] | None = None,
    ) -> tuple[CriticResult, EvaluationSummary]:
        critic_result = self.critic_agent.review(
            goal=task.goal,
            patch=patch,
            file_writes=file_writes,
            repo_context=task.repo_context,
            validation_result=validation_summary,
        )
        evaluation_summary = self.evaluator_agent.evaluate(
            goal=task.goal,
            repo_context=task.repo_context,
            validation_result=validation_summary,
            critic_result=critic_result,
        )
        return critic_result, evaluation_summary

    async def finalize_task(self, task: TaskState, *, success: bool) -> str:
        diff_text = ""
        if task.workspace_path and task.workspace_path != task.repo_path and success:
            commit_result = await self.workspace_manager.commit_snapshot(
                self._snapshot_for_task(task),
                task.repo_path,
            )
            diff_text = str(commit_result.data.get("diff", "")) if commit_result.data else commit_result.stdout
        elif task.workspace_path and task.workspace_path != task.repo_path:
            await self.workspace_manager.discard_snapshot(self._snapshot_for_task(task))

        await self.memory_service.store_task_result(
            task,
            overrides=task.constraints.model_overrides,
        )
        for signature in self._pattern_signatures(task):
            await self.memory_service.record_pattern(
                signature,
                task.summary or signature,
                overrides=task.constraints.model_overrides,
            )
        if self.rag_knowledge_base and task.rag_retrieval_events:
            self.rag_knowledge_base.mark_task_retrieval_outcome(task.task_id, success=success)
        return diff_text

    def promote_successful_fix(self, task: TaskState, *, final_diff: str) -> list[str]:
        if not self.rag_knowledge_base or not self.embedding_service or not final_diff.strip():
            return []
        try:
            examples = self.rag_knowledge_base.ingest_task_result(
                task=task,
                final_diff=final_diff,
                embedding_service=self.embedding_service,
            )
        except Exception:
            return []
        return [example.id for example in examples]

    @staticmethod
    def _snapshot_for_task(task: TaskState):
        from ai_orchestrator.schemas.workflow_models import WorkspaceSnapshotRecord

        return WorkspaceSnapshotRecord(
            task_id=task.task_id,
            source_repo_path=str(task.repo_path),
            workspace_path=str(task.workspace_path),
            strategy=task.workspace_strategy or "copy",
        )

    @staticmethod
    def _format_memory_hit(hit) -> str:
        metadata = getattr(hit, "metadata", {}) or {}
        parts = [hit.summary.strip()]
        task_mode = str(metadata.get("task_mode", "")).strip()
        if task_mode:
            parts.append(f"mode={task_mode}")
        focus_files = [str(item) for item in metadata.get("focus_files", []) if str(item).strip()]
        if focus_files:
            parts.append(f"files={', '.join(focus_files[:3])}")
        validation_failures = [str(item) for item in metadata.get("validation_failures", []) if str(item).strip()]
        if validation_failures:
            parts.append(f"failures={', '.join(validation_failures[:3])}")
        issue_patterns = [str(item) for item in metadata.get("issue_patterns", []) if str(item).strip()]
        if issue_patterns:
            parts.append(f"patterns={', '.join(issue_patterns[:3])}")
        return " | ".join(part for part in parts if part)

    @staticmethod
    def _pattern_signatures(task: TaskState) -> list[str]:
        signatures: set[str] = set()
        if task.error_message:
            signatures.add(task.error_message)
        for summary in task.validation_history:
            if summary.failed_checks:
                signatures.add("validation:" + "|".join(sorted(summary.failed_checks)))
        if task.reflection_issues:
            signatures.add("reflection:" + "|".join(sorted(set(task.reflection_issues))))
        if task.focus_issue_inventory is not None and task.focus_issue_inventory.issues:
            signatures.add(
                "focus:" + "|".join(sorted({issue.issue_type for issue in task.focus_issue_inventory.issues if issue.issue_type}))
            )
        return [signature for signature in signatures if signature.strip()]
