"""vp orchestrate subcommands: init, plan, execute, merge."""

from __future__ import annotations

import sys
import time
from pathlib import Path
from typing import Annotated, Any

import typer
import yaml

from vibepod.core.agents import agent_config_dir, effective_agent_image, get_agent_spec
from vibepod.core.config import get_config
from vibepod.core.docker import DockerClientError, DockerManager
from vibepod.core.orchestrate.plan_file import (
    OrchestratePlan,
    OrchestrateTask,
    new_plan,
    parse_plan,
    update_task_status,
    write_plan,
)
from vibepod.core.orchestrate.prompts import build_merge_prompt, build_planning_prompt
from vibepod.core.orchestrate.worker import (
    StatusResult,
    all_deps_done,
    any_dep_failed,
    read_status_file,
    tasks_ready_to_run,
)
from vibepod.core.orchestrate.worktree import (
    WorktreeError,
    create_worktree,
    delete_worktree,
    reset_worktree,
    worktree_branch_name,
)
from vibepod.utils.console import error, info, success, warning

app = typer.Typer(help="Agent-to-Agent orchestration: init, plan, execute, merge")


def _noninteractive_command(agent_name: str) -> list[str]:
    """Build a bash command that runs the agent non-interactively with VP_PROMPT_FILE.

    Uses the agent's ikwid_args and the -p flag (supported by Claude and Codex)
    to pass the prompt without user interaction.
    """
    spec = get_agent_spec(agent_name)
    flags = " ".join(spec.ikwid_args or [])
    return ["bash", "-c", f'{agent_name} {flags} -p "$(cat $VP_PROMPT_FILE)"']


def _run_controller(
    controller: str,
    workspace: Path,
    prompt_path: Path,
    config: dict,
    manager: DockerManager,
) -> bool:
    """Run the controller agent non-interactively with the given prompt file.

    Streams container output to stdout and waits for the container to exit.
    Returns True if the container exited with code 0.
    """
    spec = get_agent_spec(controller)
    image = effective_agent_image(controller, config)
    cfg_dir = agent_config_dir(controller)
    cfg_dir.mkdir(parents=True, exist_ok=True)

    env = {
        **spec.extra_env,
        "VP_PROMPT_FILE": "/vibepod-prompt.md",
    }

    container = manager.run_agent(
        agent=controller,
        image=image,
        workspace=workspace,
        config_dir=cfg_dir,
        config_mount_path=spec.config_mount_path,
        env=env,
        command=_noninteractive_command(controller),
        auto_remove=False,
        name=None,
        version="",
        network=str(config.get("network", "vibepod-network")),
        extra_volumes=[(str(prompt_path), "/vibepod-prompt.md", "ro")],
        tty=False,
        stdin_open=False,
    )

    for chunk in container.logs(stream=True, follow=True):
        sys.stdout.buffer.write(chunk)
        sys.stdout.buffer.flush()

    result = container.wait()
    try:
        container.remove()
    except Exception:
        pass

    return result.get("StatusCode", 1) == 0

_DEFAULT_PLAN_FILE = ".vibepod/PLAN.md"

_STARTER_ORCHESTRATE_CONFIG: dict[str, Any] = {
    "controller": "claude",
    "merge_strategy": "auto",
    "cleanup": True,
    "max_parallel": 4,
    "personas": {
        "frontend": {
            "preferred": "claude",
            "fallback": "codex",
            "description": "Builds UI components, handles styling, and client-side logic",
            "capabilities": [
                "React / Vue / Svelte components",
                "CSS and styling systems",
                "Accessibility (WCAG)",
                "Client-side state management",
            ],
        },
        "backend": {
            "preferred": "codex",
            "fallback": "claude",
            "description": "Implements APIs, database models, and server-side business logic",
            "capabilities": [
                "REST / GraphQL APIs",
                "Database design and migrations",
                "Authentication and authorization",
                "Performance optimization",
            ],
        },
        "security": {
            "preferred": "claude",
            "fallback": "codex",
            "description": "Reviews and implements security controls, audits dependencies",
            "capabilities": [
                "Threat modelling",
                "Dependency vulnerability scanning",
                "Auth and session security",
                "Input validation and sanitization",
            ],
        },
        "testing": {
            "preferred": "claude",
            "fallback": "codex",
            "description": "Writes and maintains automated tests at all levels",
            "capabilities": [
                "Unit and integration tests",
                "End-to-end test suites",
                "Test coverage analysis",
                "Mocking and fixtures",
            ],
        },
    },
}
_POLL_INTERVAL = 5.0
_WORKER_TIMEOUT = 3600.0


def _vibepod_dir(workspace: Path) -> Path:
    return workspace / ".vibepod"


def _status_dir(workspace: Path) -> Path:
    d = _vibepod_dir(workspace) / "status"
    d.mkdir(parents=True, exist_ok=True)
    return d


def _tasks_dir(workspace: Path) -> Path:
    d = _vibepod_dir(workspace) / "tasks"
    d.mkdir(parents=True, exist_ok=True)
    return d


def _worktrees_dir(workspace: Path) -> Path:
    return _vibepod_dir(workspace) / "worktrees"


# ---------------------------------------------------------------------------
# vp orchestrate init
# ---------------------------------------------------------------------------

@app.command("init")
def orchestrate_init(
    workspace: Annotated[
        Path, typer.Option("-w", "--workspace", help="Workspace directory")
    ] = Path("."),
    force: Annotated[
        bool, typer.Option("--force", help="Overwrite existing orchestrate config without prompting")
    ] = False,
) -> None:
    """Add starter orchestrate configuration to .vibepod/config.yaml."""
    workspace = workspace.expanduser().resolve()
    config_path = workspace / ".vibepod" / "config.yaml"
    config_path.parent.mkdir(parents=True, exist_ok=True)

    # Load existing config (or start fresh)
    existing: dict[str, Any] = {}
    if config_path.exists():
        try:
            loaded = yaml.safe_load(config_path.read_text(encoding="utf-8"))
            if isinstance(loaded, dict):
                existing = loaded
        except yaml.YAMLError as exc:
            error(f"Failed to parse {config_path}: {exc}")
            raise typer.Exit(1) from exc

    if "orchestrate" in existing and not force:
        error(f"orchestrate config already exists in {config_path}")
        error("Use --force to overwrite.")
        raise typer.Exit(1)

    existing.setdefault("version", 1)
    existing["orchestrate"] = _STARTER_ORCHESTRATE_CONFIG

    try:
        config_path.write_text(yaml.safe_dump(existing, sort_keys=False), encoding="utf-8")
    except OSError as exc:
        error(f"Failed to write {config_path}: {exc}")
        raise typer.Exit(1) from exc

    success(f"Orchestrate config written to {config_path}")
    info("")
    info("Next steps:")
    info("  1. Edit the personas in .vibepod/config.yaml to match your project and preferred agents")
    info("  2. Run: vp orchestrate plan \"<your task description>\"")
    info("  3. Review .vibepod/PLAN.md, then run: vp orchestrate execute")
    info("  4. When all tasks are done, run: vp orchestrate merge")


# ---------------------------------------------------------------------------
# vp orchestrate plan
# ---------------------------------------------------------------------------

@app.command("plan")
def orchestrate_plan(
    task: Annotated[str, typer.Argument(help="Task description for the orchestrator")],
    workspace: Annotated[
        Path, typer.Option("-w", "--workspace", help="Workspace directory")
    ] = Path("."),
    plan_file: Annotated[
        Path, typer.Option("--plan-file", help="Override PLAN.md location")
    ] = Path(_DEFAULT_PLAN_FILE),
    force: Annotated[
        bool, typer.Option("--force", help="Overwrite existing PLAN.md without prompting")
    ] = False,
) -> None:
    """Generate a PLAN.md by running the controller agent on the given task."""
    workspace = workspace.expanduser().resolve()
    plan_path = (workspace / plan_file) if not plan_file.is_absolute() else plan_file

    if plan_path.exists() and not force:
        if not typer.confirm(f"{plan_path} already exists. Overwrite?", default=False):
            raise typer.Exit(0)

    config = get_config()
    orch_cfg = config.get("orchestrate", {})
    controller = str(orch_cfg.get("controller", config.get("default_agent", "claude")))
    personas: dict = orch_cfg.get("personas", {})

    if not personas:
        error("No personas defined in config. Add orchestrate.personas to config.yaml.")
        raise typer.Exit(1)

    plan = new_plan(task)
    plan_path.parent.mkdir(parents=True, exist_ok=True)

    # Write planning prompt to .vibepod/planning-prompt.md
    planning_prompt_path = workspace / ".vibepod" / "planning-prompt.md"
    planning_prompt_path.parent.mkdir(parents=True, exist_ok=True)
    planning_prompt_path.write_text(
        build_planning_prompt(task, personas, run_id=plan.run_id),
        encoding="utf-8",
    )

    try:
        manager = DockerManager()
    except DockerClientError as exc:
        error(str(exc))
        raise typer.Exit(1) from exc

    info(f"Running controller agent ({controller}) to generate plan...")
    _run_controller(controller, workspace, planning_prompt_path, config, manager)

    if not plan_path.exists():
        error(f"Controller agent did not produce {plan_path}. Check agent output above.")
        raise typer.Exit(1)

    success(f"Plan written to {plan_path}")
    info("Review and edit the plan, then run: vp orchestrate execute")


# ---------------------------------------------------------------------------
# vp orchestrate execute
# ---------------------------------------------------------------------------

@app.command("execute")
def orchestrate_execute(
    workspace: Annotated[
        Path, typer.Option("-w", "--workspace", help="Workspace directory")
    ] = Path("."),
    plan_file: Annotated[
        Path, typer.Option("--plan-file", help="Override PLAN.md location")
    ] = Path(_DEFAULT_PLAN_FILE),
) -> None:
    """Read PLAN.md, spawn worker agents, monitor until all tasks complete."""
    workspace = workspace.expanduser().resolve()
    plan_path = (workspace / plan_file) if not plan_file.is_absolute() else plan_file

    if not plan_path.exists():
        error(f"Plan file not found: {plan_path}. Run: vp orchestrate plan <task>")
        raise typer.Exit(1)

    config = get_config()
    orch_cfg = config.get("orchestrate", {})
    max_parallel: int = int(orch_cfg.get("max_parallel", 4))

    try:
        manager = DockerManager()
    except DockerClientError as exc:
        error(str(exc))
        raise typer.Exit(1) from exc

    plan = parse_plan(plan_path.read_text(encoding="utf-8"))

    # Re-entrancy: reset running tasks whose containers are gone
    _recheck_running_tasks(plan, plan_path, manager)

    # Mark blocked tasks and waiting tasks
    _mark_blocked_tasks(plan, plan_path)
    _mark_waiting_tasks(plan, plan_path)

    status_dir = _status_dir(workspace)
    worktrees_base = _worktrees_dir(workspace)
    active: dict[str, object] = {}  # task_id -> container

    info(f"Executing plan (run-id: {plan.run_id}, max_parallel: {max_parallel})")

    while True:
        plan = parse_plan(plan_path.read_text(encoding="utf-8"))

        # Check active workers for completion
        for task_id in list(active.keys()):
            container = active[task_id]
            status_file = status_dir / f"{task_id.lower()}.json"
            result = read_status_file(status_file)
            if result is None:
                try:
                    container.reload()  # type: ignore[union-attr]
                    if container.status not in ("running", "created"):  # type: ignore[union-attr]
                        result = StatusResult(status="failed", reason="container exited without status")
                except Exception:
                    result = StatusResult(status="failed", reason="container unreachable")

            if result is not None:
                _handle_task_result(
                    task_id=task_id,
                    result=result,
                    plan=plan,
                    plan_path=plan_path,
                    workspace=workspace,
                    config=config,
                    manager=manager,
                    active=active,
                    status_dir=status_dir,
                    worktrees_base=worktrees_base,
                )
                del active[task_id]

        # Mark any newly blocked or waiting tasks
        plan = parse_plan(plan_path.read_text(encoding="utf-8"))
        _mark_blocked_tasks(plan, plan_path)
        _mark_waiting_tasks(plan, plan_path)
        plan = parse_plan(plan_path.read_text(encoding="utf-8"))

        # Spawn ready tasks (up to max_parallel)
        available_slots = max_parallel - len(active)
        if available_slots > 0:
            ready = tasks_ready_to_run(plan)
            for task in ready[:available_slots]:
                _spawn_worker(task, plan, workspace, config, manager, status_dir, worktrees_base, active, plan_path)

        # Check if all terminal
        plan = parse_plan(plan_path.read_text(encoding="utf-8"))
        terminal = {"done", "failed", "blocked"}
        if all(t.status in terminal for t in plan.tasks) and not active:
            break

        time.sleep(_POLL_INTERVAL)

    _print_execute_summary(plan)
    info("Run: vp orchestrate merge")


def _recheck_running_tasks(plan: OrchestratePlan, plan_path: Path, manager: DockerManager) -> None:
    """Reset tasks marked 'running' whose containers are no longer alive."""
    for task in plan.tasks:
        if task.status != "running":
            continue
        try:
            containers = manager.client.containers.list(
                filters={"label": f"vibepod.orchestrate.task={task.id}"}
            )
            if containers:
                continue  # still running
        except Exception:
            pass
        update_task_status(plan_path, task.id, "pending")
        warning(f"{task.id}: container not found, reset to pending")


def _mark_blocked_tasks(plan: OrchestratePlan, plan_path: Path) -> None:
    """Mark tasks as blocked if any dependency failed."""
    for task in plan.tasks:
        if task.status in ("pending", "waiting") and any_dep_failed(task, plan):
            update_task_status(plan_path, task.id, "blocked")
            warning(f"{task.id}: blocked (dependency failed)")


def _mark_waiting_tasks(plan: OrchestratePlan, plan_path: Path) -> None:
    """Mark tasks as waiting when they have unmet dependencies but no failure."""
    for task in plan.tasks:
        if task.status != "pending":
            continue
        if task.depends_on and not all_deps_done(task, plan) and not any_dep_failed(task, plan):
            update_task_status(plan_path, task.id, "waiting")


def _spawn_worker(
    task: OrchestrateTask,
    plan: OrchestratePlan,
    workspace: Path,
    config: dict,
    manager: DockerManager,
    status_dir: Path,
    worktrees_base: Path,
    active: dict,
    plan_path: Path,
) -> None:
    branch = worktree_branch_name(task.id, plan.run_id)
    worktree_path = worktrees_base / task.id.lower()

    try:
        create_worktree(workspace, worktree_path, branch)
    except WorktreeError as exc:
        error(f"{task.id}: failed to create worktree: {exc}")
        update_task_status(plan_path, task.id, "failed")
        return

    # Copy task prompt into worktree
    task_prompt_src = workspace / task.prompt_file
    prompt_dest = worktree_path / ".vibepod-task.md"
    if task_prompt_src.exists():
        prompt_dest.write_text(task_prompt_src.read_text(encoding="utf-8"), encoding="utf-8")

    status_file = status_dir / f"{task.id.lower()}.json"
    status_file.unlink(missing_ok=True)

    agent = task.agent
    orch_cfg = config.get("orchestrate", {})
    personas = orch_cfg.get("personas", {})
    persona_cfg = personas.get(task.persona, {})
    preferred = persona_cfg.get("preferred", agent)

    worker_spec = get_agent_spec(preferred)
    worker_image = effective_agent_image(preferred, config)
    worker_cfg_dir = agent_config_dir(preferred)
    worker_cfg_dir.mkdir(parents=True, exist_ok=True)

    try:
        container = manager.run_agent(
            agent=preferred,
            image=worker_image,
            workspace=worktree_path,
            config_dir=worker_cfg_dir,
            config_mount_path=worker_spec.config_mount_path,
            env={
                **worker_spec.extra_env,
                "VP_STATUS_FILE": f"/vibepod-status/{task.id.lower()}.json",
                "VP_PROMPT_FILE": "/vibepod-prompt.md",
            },
            command=_noninteractive_command(preferred),
            auto_remove=True,
            name=None,
            version="",
            network=str(config.get("network", "vibepod-network")),
            extra_volumes=[
                (str(prompt_dest), "/vibepod-prompt.md", "ro"),
                (str(status_dir), "/vibepod-status", "rw"),
            ],
            labels={"vibepod.orchestrate.task": task.id},
            tty=False,
            stdin_open=False,
        )
        update_task_status(plan_path, task.id, "running")
        active[task.id] = container
        info(f"{task.id}: started ({preferred}) in {worktree_path.name}")
    except Exception as exc:
        error(f"{task.id}: failed to start: {exc}")
        update_task_status(plan_path, task.id, "failed")


def _handle_task_result(
    task_id: str,
    result: StatusResult,
    plan: OrchestratePlan,
    plan_path: Path,
    workspace: Path,
    config: dict,
    manager: DockerManager,
    active: dict,
    status_dir: Path,
    worktrees_base: Path,
) -> None:
    task = next((t for t in plan.tasks if t.id == task_id), None)
    if task is None:
        return

    if result.status == "done":
        update_task_status(plan_path, task_id, "done")
        task.summary = result.summary  # type: ignore[attr-defined]
        success(f"{task_id}: done — {result.summary}")
    else:
        warning(f"{task_id}: failed — {result.reason}")
        if task.fallback:
            info(f"{task_id}: retrying with fallback agent ({task.fallback})...")
            worktree_path = worktrees_base / task_id.lower()
            try:
                reset_worktree(worktree_path)
            except WorktreeError as exc:
                error(f"{task_id}: could not reset worktree: {exc}")
                update_task_status(plan_path, task_id, "failed")
                return
            original_agent = task.agent
            task.agent = task.fallback
            _spawn_worker(task, plan, workspace, config, manager, status_dir, worktrees_base, active, plan_path)
            task.agent = original_agent
        else:
            update_task_status(plan_path, task_id, "failed")


def _print_execute_summary(plan: OrchestratePlan) -> None:
    done = [t for t in plan.tasks if t.status == "done"]
    failed = [t for t in plan.tasks if t.status == "failed"]
    blocked = [t for t in plan.tasks if t.status == "blocked"]
    info(f"\nExecute complete: {len(done)} done, {len(failed)} failed, {len(blocked)} blocked")
    for t in failed:
        warning(f"  FAILED: {t.id} — {t.title}")
    for t in blocked:
        warning(f"  BLOCKED: {t.id} — {t.title}")


# ---------------------------------------------------------------------------
# vp orchestrate merge
# ---------------------------------------------------------------------------

@app.command("merge")
def orchestrate_merge(
    workspace: Annotated[
        Path, typer.Option("-w", "--workspace", help="Workspace directory")
    ] = Path("."),
    plan_file: Annotated[
        Path, typer.Option("--plan-file", help="Override PLAN.md location")
    ] = Path(_DEFAULT_PLAN_FILE),
    force: Annotated[
        bool, typer.Option("--force", help="Skip warning when tasks are failed/blocked")
    ] = False,
    merge_strategy: Annotated[
        str | None,
        typer.Option("--merge-strategy", help="Override merge strategy (auto|review)")
    ] = None,
) -> None:
    """Synthesize all worker worktree results using the controller agent."""
    workspace = workspace.expanduser().resolve()
    plan_path = (workspace / plan_file) if not plan_file.is_absolute() else plan_file

    if not plan_path.exists():
        error(f"Plan file not found: {plan_path}")
        raise typer.Exit(1)

    plan = parse_plan(plan_path.read_text(encoding="utf-8"))

    failed_or_blocked = [t for t in plan.tasks if t.status in ("failed", "blocked")]
    done_tasks = [t for t in plan.tasks if t.status == "done"]

    if not done_tasks:
        error("No completed tasks to merge.")
        raise typer.Exit(1)

    if failed_or_blocked and not force:
        warning(f"{len(failed_or_blocked)} task(s) are failed/blocked:")
        for t in failed_or_blocked:
            warning(f"  {t.status.upper()}: {t.id} — {t.title}")
        if not typer.confirm("Proceed with merging completed tasks only?", default=False):
            raise typer.Exit(0)

    config = get_config()
    orch_cfg = config.get("orchestrate", {})
    controller = str(orch_cfg.get("controller", config.get("default_agent", "claude")))
    strategy = merge_strategy or str(orch_cfg.get("merge_strategy", "auto"))

    # Load summaries from status files
    status_dir = _status_dir(workspace)
    for task in done_tasks:
        sf = read_status_file(status_dir / f"{task.id.lower()}.json")
        if sf:
            task.summary = sf.summary  # type: ignore[attr-defined]

    # Review gate
    info("\nCompleted tasks:")
    for task in done_tasks:
        branch = f"orchestrate/{task.id.lower()}-{plan.run_id}"
        summary = getattr(task, "summary", "")
        info(f"  {task.id}: {task.title} — {branch}")
        if summary:
            info(f"    Summary: {summary}")

    # Write merge prompt
    merge_prompt_path = workspace / ".vibepod" / "merge-prompt.md"
    merge_prompt_path.write_text(build_merge_prompt(plan, workspace=str(workspace)), encoding="utf-8")

    info(f"\nRunning controller agent ({controller}) to merge branches...")

    try:
        manager = DockerManager()
    except DockerClientError as exc:
        error(str(exc))
        raise typer.Exit(1) from exc

    if strategy == "review":
        for task in done_tasks:
            branch = f"orchestrate/{task.id.lower()}-{plan.run_id}"
            info(f"\nAbout to merge branch: {branch}")
            if not typer.confirm("Proceed?", default=True):
                info(f"Skipping {task.id}")
                continue
            _run_merge_for_task(task, plan, workspace, controller, config, manager)
    else:
        _run_controller(controller, workspace, merge_prompt_path, config, manager)

    # Cleanup
    if bool(orch_cfg.get("cleanup", True)):
        conflict_file = workspace / ".vibepod" / "merge-conflicts.md"
        if conflict_file.exists():
            warning("Merge conflicts detected. Skipping cleanup. Resolve and re-run merge.")
        else:
            _cleanup(plan, workspace, config)

    success("Merge complete.")


def _run_merge_for_task(
    task: OrchestrateTask,
    plan: OrchestratePlan,
    workspace: Path,
    controller: str,
    config: dict,
    manager: DockerManager,
) -> None:
    branch = f"orchestrate/{task.id.lower()}-{plan.run_id}"
    prompt_path = workspace / ".vibepod" / f"merge-prompt-{task.id.lower()}.md"
    prompt_path.write_text(
        f"Merge the branch `{branch}` into the current branch.\n"
        f"Resolve any conflicts, then commit with message: merge: integrate {task.id} results\n",
        encoding="utf-8",
    )
    _run_controller(controller, workspace, prompt_path, config, manager)


def _cleanup(plan: OrchestratePlan, workspace: Path, config: dict) -> None:
    worktrees_base = _worktrees_dir(workspace)
    for task in plan.tasks:
        if task.status != "done":
            continue
        branch = worktree_branch_name(task.id, plan.run_id)
        worktree_path = worktrees_base / task.id.lower()
        try:
            delete_worktree(workspace, worktree_path, branch)
        except Exception as exc:
            warning(f"Could not clean up {task.id}: {exc}")

    status_dir = _status_dir(workspace)
    for f in status_dir.glob("*.json"):
        f.unlink(missing_ok=True)

    info("Cleaned up worktrees and status files.")
