from __future__ import annotations

import asyncio
from uuid import uuid4

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.config.settings import PROJECT_ROOT, Settings
from ai_orchestrator.orchestrators.base import OrchestratorEngine
from ai_orchestrator.orchestrators.langgraph_engine import LangGraphEngine
from ai_orchestrator.orchestrators.microsoft_agent_engine import MicrosoftAgentEngine
from ai_orchestrator.runtime.docker_sandbox import DockerSandbox
from ai_orchestrator.schemas.task_request import TaskRequest
from ai_orchestrator.schemas.task_result import TaskResultResponse
from ai_orchestrator.schemas.task_state import TaskState
from ai_orchestrator.schemas.task_status import (
    TaskLifecycleStatus,
    TaskStatusResponse,
    TaskSubmittedResponse,
)
from ai_orchestrator.services.agent_controller import AgentController
from ai_orchestrator.services.artifact_store import ArtifactStore
from ai_orchestrator.services.memory_service import MemoryService
from ai_orchestrator.services.model_router import ModelRouter
from ai_orchestrator.services.postgres_store import PostgresQueueStore, PostgresTaskStore
from ai_orchestrator.services.prompt_loader import PromptLoader
from ai_orchestrator.services.rag_knowledge_base import RAGKnowledgeBase
from ai_orchestrator.services.queue_store import SQLiteQueueStore
from ai_orchestrator.services.repo_analyzer import RepoAnalyzer
from ai_orchestrator.services.safety_guardrails import SafetyGuardrails
from ai_orchestrator.services.sqlite_store import SQLiteTaskStore
from ai_orchestrator.services.supervisor_agent import SupervisorAgent
from ai_orchestrator.services.task_mode_router import TaskModeRouter
from ai_orchestrator.services.task_worker import TaskWorker
from ai_orchestrator.services.validation_engine import ValidationEngine
from ai_orchestrator.services.workspace_manager import WorkspaceManager
from ai_orchestrator.services.embedding_service import EmbeddingService
from ai_orchestrator.tools.file_tools import FileTools
from ai_orchestrator.tools.git_runner import GitRunner
from ai_orchestrator.tools.shell_runner import ShellRunner
from ai_orchestrator.tools.test_runner import TestRunner


class TaskService:
    def __init__(
        self,
        *,
        settings: Settings,
        store,
        queue,
        task_mode_router: TaskModeRouter,
        orchestrators: dict[str, OrchestratorEngine],
    ) -> None:
        self.settings = settings
        self.store = store
        self.queue = queue
        self.task_mode_router = task_mode_router
        self.orchestrators = orchestrators

    async def submit_task(self, request: TaskRequest) -> TaskSubmittedResponse:
        repo_path = request.repo_path.resolve()
        if not repo_path.exists() or not repo_path.is_dir():
            raise ValueError(f"repo_path must point to an existing directory: {repo_path}")

        engine_name = request.constraints.engine or self.settings.default_engine
        if engine_name not in self.orchestrators:
            raise ValueError(f"Unknown orchestrator engine: {engine_name}")

        task_mode = request.constraints.task_mode or self.task_mode_router.detect_mode(
            goal=request.goal,
            repo_path=repo_path,
            test_command=request.constraints.test_command,
        )
        constraints = request.constraints
        if "max_attempts" not in constraints.model_fields_set:
            constraints = constraints.model_copy(update={"max_attempts": self.settings.default_max_attempts})
        constraints = constraints.model_copy(update={"task_mode": task_mode})
        request = request.model_copy(update={"repo_path": repo_path})
        task_id = f"task_{uuid4().hex[:12]}"
        task_state = TaskState(
            task_id=task_id,
            goal=request.goal,
            repo_path=repo_path,
            task_mode=task_mode,
            constraints=constraints,
        )
        self.store.create(task_state)
        self.queue.enqueue(
            task_id=task_id,
            engine_name=engine_name,
            max_attempts=constraints.max_attempts,
        )
        return TaskSubmittedResponse(task_id=task_id, status=task_state.status)

    async def get_status(self, task_id: str) -> TaskStatusResponse | None:
        task = self.store.get(task_id)
        if task is None:
            return None
        return TaskStatusResponse(
            task_id=task.task_id,
            status=task.status,
            current_step=task.current_step,
            attempt=task.attempt,
            updated_at=task.timestamps.updated_at,
        )

    async def get_result(self, task_id: str) -> TaskResultResponse | None:
        task = self.store.get(task_id)
        if task is None:
            return None
        if task.status not in TaskLifecycleStatus.terminal():
            raise RuntimeError("Task result is not available until the task reaches a terminal state.")
        return TaskResultResponse(
            task_id=task.task_id,
            status=task.status,
            summary=task.summary,
            artifacts=task.artifacts,
            error_message=task.error_message,
        )

    async def cancel_task(self, task_id: str) -> TaskStatusResponse | None:
        task = self.store.get(task_id)
        if task is None:
            return None
        if task.status not in TaskLifecycleStatus.terminal():
            task.mark_terminal(TaskLifecycleStatus.cancelled, summary="Task cancelled by user.")
            self.store.save(task)
        self.queue.cancel_for_task(task_id)
        return await self.get_status(task_id)

    async def recover_stale_tasks(self) -> None:
        self.queue.requeue_expired()

    async def shutdown(self) -> None:
        return None

    async def wait_for_completion(self, task_id: str, timeout_seconds: float = 10.0) -> TaskState:
        deadline = asyncio.get_running_loop().time() + timeout_seconds
        while True:
            task = self.store.get(task_id)
            if task is None:
                raise LookupError(f"Unknown task: {task_id}")
            if task.status in TaskLifecycleStatus.terminal():
                return task
            if asyncio.get_running_loop().time() >= deadline:
                raise TimeoutError(f"Timed out waiting for task {task_id}")
            await asyncio.sleep(0.05)


def build_runtime_components(settings: Settings | None = None) -> dict:
    settings = settings or Settings()
    settings.prepare_dirs()

    if settings.queue_backend == "postgres":
        if not settings.postgres_dsn:
            raise RuntimeError("AIIQ_POSTGRES_DSN must be configured when queue_backend=postgres.")
        store = PostgresTaskStore(settings.postgres_dsn)
        queue = PostgresQueueStore(settings.postgres_dsn)
    else:
        store = SQLiteTaskStore(settings.database_path)
        queue = SQLiteQueueStore(settings.database_path)

    artifact_store = ArtifactStore(settings.artifacts_dir, PROJECT_ROOT)
    prompt_loader = PromptLoader(settings.prompt_dir)
    model_router = ModelRouter(settings)
    safety_guardrails = SafetyGuardrails()
    sandbox = DockerSandbox(settings)
    shell_runner = ShellRunner(sandbox, safety_guardrails=safety_guardrails)
    file_tools = FileTools(settings)
    git_runner = GitRunner(shell_runner)
    test_runner = TestRunner(shell_runner)
    memory_service = MemoryService(settings, model_router)
    repo_analyzer = RepoAnalyzer(settings)
    task_mode_router = TaskModeRouter()
    rag_knowledge_base = RAGKnowledgeBase(settings)
    embedding_service = EmbeddingService(settings)
    if settings.rag_auto_import_seed_examples and settings.rag_seed_examples_path:
        rag_knowledge_base.import_examples_file(settings.rag_seed_examples_path, embedding_service=embedding_service)
    workspace_manager = WorkspaceManager(
        settings=settings,
        shell_runner=shell_runner,
        git_runner=git_runner,
    )
    validation_engine = ValidationEngine(
        settings=settings,
        shell_runner=shell_runner,
        test_runner=test_runner,
        safety_guardrails=safety_guardrails,
        rag_knowledge_base=rag_knowledge_base,
        embedding_service=embedding_service,
    )
    planner_agent = PlannerAgent(model_router)
    critic_agent = CriticAgent()
    evaluator_agent = EvaluatorAgent()
    supervisor_agent = SupervisorAgent()
    agent_controller = AgentController(
        planner_agent=planner_agent,
        critic_agent=critic_agent,
        evaluator_agent=evaluator_agent,
        supervisor_agent=supervisor_agent,
        memory_service=memory_service,
        repo_analyzer=repo_analyzer,
        validation_engine=validation_engine,
        workspace_manager=workspace_manager,
        rag_knowledge_base=rag_knowledge_base,
        embedding_service=embedding_service,
    )

    orchestrators: dict[str, OrchestratorEngine] = {
        "langgraph": LangGraphEngine(
            settings=settings,
            store=store,
            artifact_store=artifact_store,
            prompt_loader=prompt_loader,
            model_router=model_router,
            file_tools=file_tools,
            git_runner=git_runner,
            agent_controller=agent_controller,
            rag_knowledge_base=rag_knowledge_base,
            embedding_service=embedding_service,
        ),
        "microsoft": MicrosoftAgentEngine(store),
    }

    return {
        "settings": settings,
        "store": store,
        "queue": queue,
        "task_mode_router": task_mode_router,
        "orchestrators": orchestrators,
    }


def build_task_service(settings: Settings | None = None) -> TaskService:
    runtime = build_runtime_components(settings)
    return TaskService(
        settings=runtime["settings"],
        store=runtime["store"],
        queue=runtime["queue"],
        task_mode_router=runtime["task_mode_router"],
        orchestrators=runtime["orchestrators"],
    )


def build_task_worker(settings: Settings | None = None) -> TaskWorker:
    runtime = build_runtime_components(settings)
    return TaskWorker(
        settings=runtime["settings"],
        store=runtime["store"],
        queue=runtime["queue"],
        orchestrators=runtime["orchestrators"],
    )
