"""Child process dispatch, monitoring, and harvesting.

Spawns ``claude --print`` subprocesses for each child, monitors their progress,
enforces timeouts, and harvests results (exit code, cost, task status).
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import os
import signal
from datetime import UTC, datetime
from pathlib import Path

from cleave.orchestrator.config import OrchestratorConfig
from cleave.orchestrator.errors import TimeoutExceededError
from cleave.orchestrator.state import ChildState

logger = logging.getLogger(__name__)


def _build_claude_args(
    claude_cli: str,
    prompt: str,
    config: OrchestratorConfig,
    cwd: Path,
) -> list[str]:
    """Build the argument list for a claude --print subprocess.

    Isolation flags ensure children run in a controlled environment:
    - --no-session-persistence: prevent .claude/ directory pollution in worktrees
    - --disable-slash-commands: prevent children from invoking skills (e.g. /cleave)
    - --strict-mcp-config: prevent ambient MCP server inheritance from project configs
    """
    args = [
        claude_cli,
        "--print",
        "--output-format", "json",
        "--model", config.model,
        "--max-turns", "200",
        "--permission-mode", config.permission_mode,
        "--no-session-persistence",
        "--disable-slash-commands",
    ]
    if config.child_budget_usd > 0:
        args.extend(["--max-budget-usd", str(config.child_budget_usd)])
    if config.allowed_tools:
        # --allowedTools accepts comma-separated or multiple args.
        # Use comma-separated in a single argument for safety with subprocess exec.
        tools_csv = ",".join(config.allowed_tools.split())
        args.extend(["--allowedTools", tools_csv])

    # MCP isolation: --strict-mcp-config prevents children from inheriting
    # MCP servers from project .mcp.json or settings.local.json files in the
    # directory ancestry. If explicit mcp_config is provided, pass it through;
    # otherwise just use --strict-mcp-config alone to block ambient MCP.
    if config.mcp_config:
        args.extend(["--mcp-config", config.mcp_config])
    args.append("--strict-mcp-config")

    return args


# Allowlisted environment variables safe to pass to child subprocesses.
# Everything else is stripped to prevent credential leakage.
_ENV_ALLOWLIST_PREFIXES = (
    "PATH",
    "HOME",
    "USER",
    "LOGNAME",
    "SHELL",
    "TERM",
    "LANG",
    "LC_",
    "TMPDIR",
    "TEMP",
    "TMP",
    "XDG_",
    "EDITOR",
    "VISUAL",
    # Claude CLI needs these to function
    "ANTHROPIC_",
    # Git operation support
    "GIT_",
    "SSH_AUTH_SOCK",
    # Node.js runtime (Claude CLI is Node-based)
    "NODE_",
    "NVM_",
    "FNM_",
    # Nix support
    "NIX_",
    # System
    "DISPLAY",
    "COLORTERM",
    "NO_COLOR",
)


def _build_child_env(
    child: ChildState,
    workspace_path: Path | None = None,
    task_file_path: Path | None = None,
    depth: int = 0,
    run_id: str = "",
    repo_path: Path | None = None,
) -> dict[str, str]:
    """Build a minimal environment for child subprocesses.

    Uses an allowlist approach — only known-safe variables are inherited.
    This prevents credential leakage (AWS_*, GITHUB_TOKEN, etc.) to
    children running with bypassPermissions + Bash access.

    If repo_path has a .venv, sets VIRTUAL_ENV and prepends .venv/bin
    to PATH so bare ``python``/``pytest`` commands resolve correctly
    in worktree children.
    """
    env: dict[str, str] = {}

    # Only inherit allowlisted variables
    for key, value in os.environ.items():
        if any(key.startswith(prefix) for prefix in _ENV_ALLOWLIST_PREFIXES):
            env[key] = value

    # Never pass CLAUDECODE — it blocks subprocess Claude invocation
    env.pop("CLAUDECODE", None)

    # Inject repo venv into PATH so children can use bare python/pytest
    if repo_path:
        venv_bin = repo_path / ".venv" / "bin"
        if venv_bin.is_dir():
            env["VIRTUAL_ENV"] = str(repo_path / ".venv")
            current_path = env.get("PATH", "")
            env["PATH"] = f"{venv_bin}{os.pathsep}{current_path}" if current_path else str(venv_bin)

    # Orchestrator context — helps children understand who they are
    env["CLEAVE_RUN_ID"] = run_id
    env["CLEAVE_CHILD_ID"] = str(child.child_id)
    env["CLEAVE_DEPTH"] = str(depth)
    if workspace_path:
        env["CLEAVE_WORKSPACE"] = str(workspace_path)
    if task_file_path:
        env["CLEAVE_TASK_FILE"] = str(task_file_path)

    return env


async def spawn_child(
    claude_cli: str,
    prompt: str,
    config: OrchestratorConfig,
    child: ChildState,
    cwd: Path,
    env: dict[str, str] | None = None,
) -> asyncio.subprocess.Process:
    """Spawn a claude --print subprocess for a child task.

    The prompt is passed via stdin to avoid shell escaping issues.
    """
    args = _build_claude_args(claude_cli, prompt, config, cwd)

    logger.info(
        "Spawning child %d (%s) in %s [budget=$%.2f, timeout=%ds]",
        child.child_id, child.label, cwd,
        config.child_budget_usd, config.child_timeout_seconds,
    )
    if config.verbose:
        logger.debug("Command: %s", " ".join(args))

    proc = await asyncio.create_subprocess_exec(
        *args,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        cwd=str(cwd),
        env=env,
        # Create new session so SIGTERM/SIGKILL can target the entire
        # process group (claude CLI + any child processes it spawns).
        start_new_session=True,
    )

    # Write prompt to stdin and close
    if proc.stdin is not None:
        proc.stdin.write(prompt.encode("utf-8"))
        await proc.stdin.drain()
        proc.stdin.close()

    child.status = "running"
    child.pid = proc.pid
    child.started_at = datetime.now(UTC).isoformat()
    return proc


async def wait_for_child(
    proc: asyncio.subprocess.Process,
    child: ChildState,
    timeout_seconds: int,
) -> tuple[int, str, str]:
    """Wait for a child process to complete, enforcing timeout.

    Returns (exit_code, stdout, stderr).
    On timeout, sends SIGTERM then SIGKILL.
    """
    try:
        stdout_bytes, stderr_bytes = await asyncio.wait_for(
            proc.communicate(),
            timeout=timeout_seconds,
        )
        stdout = stdout_bytes.decode("utf-8", errors="replace") if stdout_bytes else ""
        stderr = stderr_bytes.decode("utf-8", errors="replace") if stderr_bytes else ""
        return proc.returncode or 0, stdout, stderr
    except TimeoutError as exc:
        logger.warning(
            "Child %d (%s) timed out after %ds, sending SIGTERM",
            child.child_id, child.label, timeout_seconds,
        )
        partial_stderr = ""
        try:
            proc.terminate()
            # Give it 10 seconds to clean up
            try:
                _, partial_stderr_bytes = await asyncio.wait_for(
                    proc.communicate(), timeout=10,
                )
                partial_stderr = (
                    partial_stderr_bytes.decode("utf-8", errors="replace")
                    if partial_stderr_bytes else ""
                )
                if partial_stderr:
                    logger.warning(
                        "Child %d stderr on timeout:\n%s",
                        child.child_id, partial_stderr[:2000],
                    )
            except TimeoutError:
                logger.warning("Child %d did not terminate, sending SIGKILL", child.child_id)
                # Kill the entire process group (start_new_session=True)
                try:
                    os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
                except (ProcessLookupError, PermissionError):
                    proc.kill()
                # Bounded wait to avoid hanging on pipe drain
                with contextlib.suppress(TimeoutError):
                    await asyncio.wait_for(proc.communicate(), timeout=5)
        except ProcessLookupError:
            pass  # Already exited
        raise TimeoutExceededError(
            elapsed=float(timeout_seconds),
            limit=float(timeout_seconds),
            message=f"Child {child.child_id} ({child.label}) timed out after {timeout_seconds}s",
        ) from exc


def parse_claude_output(stdout: str) -> dict:
    """Parse the JSON output from claude --print --output-format json.

    Returns a dict with at minimum 'cost_usd' and 'result' keys.
    """
    if not stdout.strip():
        return {"cost_usd": 0.0, "result": ""}

    # claude --print --output-format json outputs one JSON object per line
    # or a single JSON object. We want the last complete JSON object.
    result: dict = {"cost_usd": 0.0, "result": ""}

    # Try parsing the entire output as JSON first
    try:
        data = json.loads(stdout)
        if isinstance(data, dict):
            result["cost_usd"] = data.get("cost_usd", 0.0)
            result["result"] = data.get("result", data.get("content", ""))
            result["model"] = data.get("model", "")
            result["duration_ms"] = data.get("duration_ms", 0)
            result["num_turns"] = data.get("num_turns", 0)
            return result
        elif isinstance(data, list):
            # Array of message objects — extract the last assistant message
            for msg in reversed(data):
                if isinstance(msg, dict) and msg.get("role") == "assistant":
                    result["result"] = msg.get("content", "")
                    break
            return result
    except json.JSONDecodeError:
        pass

    # Try parsing line by line (streaming JSON lines)
    for line in reversed(stdout.splitlines()):
        line = line.strip()
        if not line:
            continue
        try:
            data = json.loads(line)
            if isinstance(data, dict):
                result["cost_usd"] = data.get("cost_usd", result["cost_usd"])
                result["result"] = data.get("result", data.get("content", result["result"]))
                break
        except json.JSONDecodeError:
            continue

    # If we still don't have a result, use raw stdout
    if not result["result"]:
        result["result"] = stdout

    return result


def harvest_child_result(
    child: ChildState,
    exit_code: int,
    stdout: str,
    stderr: str,
) -> dict:
    """Process the result of a completed child.

    Updates the ChildState in-place and returns parsed output.
    """
    child.exit_code = exit_code
    child.completed_at = datetime.now(UTC).isoformat()

    parsed = parse_claude_output(stdout)
    child.cost_usd = parsed.get("cost_usd", 0.0)

    if exit_code == 0:
        child.status = "completed"
    else:
        child.status = "failed"
        child.error_message = stderr[:2000] if stderr else f"Exit code {exit_code}"

    logger.info(
        "Child %d (%s) finished: status=%s, exit_code=%d, cost=$%.4f",
        child.child_id, child.label, child.status, exit_code, child.cost_usd or 0.0,
    )

    return parsed


def _compute_dispatch_waves(children: list[ChildState]) -> list[list[ChildState]]:
    """Group children into sequential waves using topological ordering.

    Children with no unmet dependencies go in wave 0, children whose deps
    are all in earlier waves go in the next wave, etc.

    If a deadlock is detected (no children ready but some remain), all
    remaining children are dispatched in one final wave.
    """
    if not children:
        return []

    label_to_child = {c.label: c for c in children}
    remaining = set(c.label for c in children)
    satisfied: set[str] = set()
    waves: list[list[ChildState]] = []

    while remaining:
        # Find children whose deps are all satisfied (or have no deps)
        ready = []
        for label in sorted(remaining):
            child = label_to_child[label]
            deps = set(child.depends_on) & set(label_to_child.keys())
            if deps <= satisfied:
                ready.append(child)

        if not ready:
            # Deadlock breaker: dispatch all remaining
            logger.warning(
                "Dependency deadlock detected — dispatching remaining %d children",
                len(remaining),
            )
            ready = [label_to_child[label] for label in sorted(remaining)]

        waves.append(ready)
        for child in ready:
            remaining.discard(child.label)
            satisfied.add(child.label)

    return waves


async def dispatch_children(
    claude_cli: str,
    children: list[ChildState],
    prompts: dict[int, str],
    worktree_paths: dict[int, Path],
    config: OrchestratorConfig,
    total_cost_so_far: float = 0.0,
    workspace_path: Path | None = None,
    task_file_paths: dict[int, Path] | None = None,
    depth: int = 0,
    run_id: str = "",
    repo_path: Path | None = None,
) -> list[dict]:
    """Dispatch children in dependency-ordered waves.

    Children within a wave run in parallel (up to max_parallel_children).
    Waves are executed sequentially — the next wave starts only after
    all children in the current wave have finished.

    Returns list of harvest results in child_id order.
    """
    semaphore = asyncio.Semaphore(config.max_parallel_children)
    results: dict[int, dict] = {}
    # Lock protects the budget check + commit sequence so concurrent
    # children can't all pass the check before any costs are recorded.
    budget_lock = asyncio.Lock()
    committed_cost = total_cost_so_far

    async def _run_one(child: ChildState) -> None:
        nonlocal committed_cost
        try:
            async with semaphore:
                # Budget check under lock to prevent concurrent overrun
                async with budget_lock:
                    if committed_cost >= config.max_budget_usd:
                        child.status = "failed"
                        child.error_message = "Budget exhausted before dispatch"
                        results[child.child_id] = {"cost_usd": 0.0, "result": ""}
                        return
                    # Reserve this child's budget slot
                    committed_cost += config.child_budget_usd

                cwd = worktree_paths.get(child.child_id, config.repo_path)
                prompt = prompts[child.child_id]
                task_file = task_file_paths.get(child.child_id) if task_file_paths else None

                child_env = _build_child_env(
                    child,
                    workspace_path=workspace_path,
                    task_file_path=task_file,
                    depth=depth,
                    run_id=run_id,
                    repo_path=repo_path,
                )

                proc = await spawn_child(claude_cli, prompt, config, child, cwd, env=child_env)
                try:
                    exit_code, stdout, stderr = await wait_for_child(
                        proc, child, config.child_timeout_seconds
                    )
                except TimeoutExceededError:
                    child.status = "failed"
                    child.error_message = f"Timed out after {config.child_timeout_seconds}s"
                    results[child.child_id] = {"cost_usd": 0.0, "result": ""}
                    return

                parsed = harvest_child_result(child, exit_code, stdout, stderr)
                results[child.child_id] = parsed

                # Reconcile actual cost vs reserved budget
                actual_cost = parsed.get("cost_usd", 0.0)
                async with budget_lock:
                    committed_cost += actual_cost - config.child_budget_usd

        except Exception as exc:
            # Catch-all: ensure child is marked failed, not left as "pending"
            logger.error(
                "Child %d (%s) failed with unexpected error: %s",
                child.child_id, child.label, exc,
            )
            if child.status not in ("completed", "failed"):
                child.status = "failed"
                child.error_message = f"Unexpected error: {exc}"
            results[child.child_id] = {"cost_usd": 0.0, "result": ""}

    # Only dispatch pending children
    pending = [c for c in children if c.status == "pending"]
    waves = _compute_dispatch_waves(pending)

    for i, wave in enumerate(waves):
        logger.info(
            "Dispatching wave %d/%d: %s",
            i + 1, len(waves), [c.label for c in wave],
        )
        tasks = [asyncio.create_task(_run_one(child)) for child in wave]
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    # Return in child_id order
    return [results.get(c.child_id, {"cost_usd": 0.0, "result": ""}) for c in children]


async def kill_all_children(children: list[ChildState]) -> None:
    """Send SIGTERM to all running children's process groups.

    Uses os.killpg to terminate the entire process group (since children
    are spawned with start_new_session=True). Falls back to os.kill if
    process group lookup fails.
    """
    for child in children:
        if child.status == "running" and child.pid:
            try:
                pgid = os.getpgid(child.pid)
                os.killpg(pgid, signal.SIGTERM)
                logger.info(
                    "Sent SIGTERM to child %d process group (pid %d, pgid %d)",
                    child.child_id, child.pid, pgid,
                )
            except (ProcessLookupError, PermissionError):
                pass  # Already exited or PID reused
            except OSError:
                # Fallback: try direct PID signal
                with contextlib.suppress(ProcessLookupError, PermissionError):
                    os.kill(child.pid, signal.SIGTERM)
