"""Dispatcher: spawns Claude Code CLI subprocesses to execute Dobby tasks."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import os
import signal
from asyncio.subprocess import PIPE, STDOUT
from datetime import UTC, datetime

from dobby.core.collector import Collector
from dobby.core.prompt_builder import PromptBuilder
from dobby.core.worktree import cleanup_worktree, create_worktree, merge_worktree
from dobby.models import Project, Result, Task

logger = logging.getLogger(__name__)


class Dispatcher:
    """Spawns and manages Claude Code agent subprocesses.

    Concurrency is controlled via an asyncio.Semaphore so that no more than
    ``max_concurrent`` agents run at the same time.
    """

    def __init__(
        self,
        max_concurrent: int,
        default_timeout: int,
        results_dir: str,
        store=None,
    ) -> None:
        self.max_concurrent = max_concurrent
        self.default_timeout = default_timeout
        self.results_dir = results_dir
        self._store = store

        self._semaphore: asyncio.Semaphore = asyncio.Semaphore(max_concurrent)
        # Maps task_id -> subprocess PID for in-flight agents
        self.running_pids: dict[str, int] = {}

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def active_count(self) -> int:
        """Number of agents currently running (semaphore slots consumed)."""
        return self.max_concurrent - self._semaphore._value  # type: ignore[attr-defined]

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def dispatch(
        self,
        task: Task,
        project: Project,
        timeout: int | None = None,
        infra_name: str | None = None,
        user_preferences: list[dict] | None = None,
    ) -> Result:
        """Dispatch a task to a Claude Code subprocess.

        If the project's dispatch_to is "remote:<host>", the agent runs on
        the remote host via SSH. Otherwise it runs locally.
        """
        if project.is_remote:
            return await self._dispatch_remote(task, project, timeout, infra_name, user_preferences)
        return await self._dispatch_local(task, project, timeout, infra_name, user_preferences)

    async def _dispatch_local(
        self,
        task: Task,
        project: Project,
        timeout: int | None = None,
        infra_name: str | None = None,
        user_preferences: list[dict] | None = None,
    ) -> Result:
        """Dispatch a task to a local Claude Code subprocess.

        Uses git worktrees for parallel execution — each task gets its own
        worktree branch, runs there, then merges back to main on completion.
        """
        effective_timeout = timeout if timeout is not None else self.default_timeout
        started_at = datetime.now(UTC)

        # 1. Build prompt
        prompt = PromptBuilder.build(task, project, infra_name, user_preferences=user_preferences)

        # 2. Create worktree for parallel execution (unless resuming a session)
        worktree_path = None
        use_worktree = not task.session_id  # don't use worktree when resuming
        if use_worktree:
            worktree_path = await create_worktree(project.repo, task.id)
            if worktree_path:
                logger.info("Task %s running in worktree: %s", task.id, worktree_path)

        cwd = worktree_path or project.repo

        # 3. Acquire semaphore (blocks until a slot is free)
        async with self._semaphore:
            # 4. Build command based on backend
            if task.backend == "codex":
                args = [
                    "codex", "exec",
                    "--json",
                    "-s", "danger-full-access",
                    "-c", 'model_reasoning_effort="high"',
                    "-C", cwd,
                    prompt,
                ]
            else:
                # Default: Claude Code
                args = [
                    "claude",
                    "-p",
                    prompt,
                    "--output-format",
                    "stream-json",
                    "--verbose",
                    "--max-turns",
                    "200",
                    "--dangerously-skip-permissions",
                ]
                # Resume an existing session if session_id is set
                if task.session_id:
                    args.extend(["--resume", task.session_id])

            # 5. Ensure results directory exists for this task
            task_results_dir = os.path.join(self.results_dir, task.id)
            os.makedirs(task_results_dir, exist_ok=True)

            # 5b. Create execution span for trajectory tracking
            if task.trace_id and self._store:
                try:
                    from dobby.core.trajectory import trajectory_parser
                    log_path_for_span = os.path.join(task_results_dir, "agent.log")
                    await trajectory_parser.add_span(
                        self._store, task.trace_id, kind="execution",
                        task_id=task.id, log_path=log_path_for_span,
                        metadata={"dispatch": "local"},
                    )
                except Exception:
                    pass

            # 6. Spawn subprocess in the worktree (or main repo)
            proc = await asyncio.create_subprocess_exec(
                *args,
                stdout=PIPE,
                stderr=STDOUT,
                cwd=cwd,
                limit=10 * 1024 * 1024,
            )

            # 7. Track PID
            self.running_pids[task.id] = proc.pid

            # 8. Stream stdout to log file in real-time
            log_path = os.path.join(task_results_dir, "agent.log")
            output_lines: list[str] = []

            try:
                with open(log_path, "w") as log_file:
                    while True:
                        line = await asyncio.wait_for(
                            proc.stdout.readline(),
                            timeout=effective_timeout,
                        )
                        if not line:
                            break
                        decoded = line.decode(errors="replace")
                        output_lines.append(decoded)
                        log_file.write(decoded)
                        log_file.flush()

            except TimeoutError:
                with contextlib.suppress(ProcessLookupError):
                    proc.kill()
                # Clean up worktree on timeout
                if worktree_path:
                    await cleanup_worktree(project.repo, task.id)
                return Result(
                    task_id=task.id,
                    status="failed",
                    summary=f"Task timed out after {effective_timeout}s",
                    started_at=started_at,
                )
            finally:
                self.running_pids.pop(task.id, None)

            await proc.wait()
            output = "".join(output_lines)
            exit_code = proc.returncode or 0

            # 9. Parse output
            result = Collector.parse_output(output, exit_code, task.id)
            result.started_at = started_at

            # 10. Save log artifacts
            log_path = Collector.save_artifacts(task.id, output, self.results_dir)
            result.artifacts.append({"type": "log", "path": log_path})

            # 11. Merge worktree back to main
            if worktree_path and result.status == "completed":
                merge_ok, merge_msg = await merge_worktree(project.repo, task.id)
                if merge_ok:
                    logger.info("Task %s: %s", task.id, merge_msg)
                    result.metrics["merge"] = merge_msg
                else:
                    logger.warning("Task %s merge failed: %s", task.id, merge_msg)
                    result.metrics["merge"] = merge_msg
                    result.summary = (result.summary or "") + f" (merge: {merge_msg})"
            elif worktree_path:
                # Task failed — clean up worktree without merging
                await cleanup_worktree(project.repo, task.id)

            return result

    async def _dispatch_remote(
        self,
        task: Task,
        project: Project,
        timeout: int | None = None,
        infra_name: str | None = None,
        user_preferences: list[dict] | None = None,
    ) -> Result:
        """Dispatch a task to a remote host via SSH."""
        effective_timeout = timeout if timeout is not None else self.default_timeout
        started_at = datetime.now(UTC)
        host = project.remote_host

        # Build prompt
        prompt = PromptBuilder.build(task, project, infra_name, user_preferences=user_preferences)

        async with self._semaphore:
            # Build the remote command — use the repo path on the remote host
            escaped_prompt = prompt.replace("'", "'\\''")
            if task.backend == "codex":
                remote_cmd = (
                    f"cd {project.repo} && "
                    f"codex exec --json -s danger-full-access "
                    f"-c 'model_reasoning_effort=\"high\"' "
                    f"-C {project.repo} '{escaped_prompt}'"
                )
            else:
                remote_cmd = (
                    f"cd {project.repo} && "
                    f"claude -p '{escaped_prompt}' "
                    f"--output-format json --max-turns 200 "
                    f"--dangerously-skip-permissions"
                )
                if task.session_id:
                    remote_cmd += f" --resume {task.session_id}"

            ssh_args = [
                "ssh",
                "-o", "StrictHostKeyChecking=no",
                "-o", "UserKnownHostsFile=/dev/null",
                "-o", "BatchMode=yes",
                f"root@{host}",
                remote_cmd,
            ]

            # Ensure results directory exists
            task_results_dir = os.path.join(self.results_dir, task.id)
            os.makedirs(task_results_dir, exist_ok=True)

            # Spawn SSH subprocess
            proc = await asyncio.create_subprocess_exec(
                *ssh_args,
                stdout=PIPE,
                stderr=STDOUT,
            )

            self.running_pids[task.id] = proc.pid

            log_path = os.path.join(task_results_dir, "agent.log")
            output_lines: list[str] = []

            try:
                with open(log_path, "w") as log_file:
                    while True:
                        line = await asyncio.wait_for(
                            proc.stdout.readline(),
                            timeout=effective_timeout,
                        )
                        if not line:
                            break
                        decoded = line.decode(errors="replace")
                        output_lines.append(decoded)
                        log_file.write(decoded)
                        log_file.flush()
            except TimeoutError:
                with contextlib.suppress(ProcessLookupError):
                    proc.kill()
                return Result(
                    task_id=task.id,
                    status="failed",
                    summary=f"Remote task timed out after {effective_timeout}s",
                    started_at=started_at,
                )
            finally:
                self.running_pids.pop(task.id, None)

            await proc.wait()
            output = "".join(output_lines)
            exit_code = proc.returncode or 0

            result = Collector.parse_output(output, exit_code, task.id)
            result.started_at = started_at

            log_path = Collector.save_artifacts(task.id, output, self.results_dir)
            result.artifacts.append({"type": "log", "path": log_path})

            return result

    async def cancel(self, task_id: str) -> None:
        """Cancel a running task by sending SIGTERM, then SIGKILL after 10s."""
        pid = self.running_pids.get(task_id)
        if pid is None:
            return

        try:
            os.kill(pid, signal.SIGTERM)
        except ProcessLookupError:
            return

        # Wait up to 10 seconds for graceful shutdown
        for _ in range(100):
            await asyncio.sleep(0.1)
            try:
                os.kill(pid, 0)  # check if still alive
            except ProcessLookupError:
                return  # process exited cleanly

        # Force kill if still alive
        with contextlib.suppress(ProcessLookupError):
            os.kill(pid, signal.SIGKILL)
