from __future__ import annotations

import shlex
import subprocess
from pathlib import Path
from typing import Iterable, List, Optional, Tuple

import click
import questionary

from ..agents import configured_agent_vms, find_agent
from ..agents_modules import get_agent_class
from ..ansible_utils import (
    AnsibleCollectionError,
    ansible_playbook_command,
    count_playbook_tasks,
    ensure_multipass_collection,
    run_ansible_playbook,
)
from ..config import AgentConfig, ConfigError, VmConfig, load_agents_config, load_config, load_vms_config, resolve_config_path
from ..debug import debug_scope
from ..interactive import is_interactive_terminal
from ..i18n import tr
from ..progress import ProgressManager
from ..vm import MultipassError, ensure_multipass_available, resolve_proxychains
from . import debug_option, non_interactive_option

PLAYBOOKS_DIR = Path(__file__).resolve().parents[1] / "ansible" / "agents"
_ALL_AGENTS_VALUE = "__all_agents__"
_DEFAULT_VM_VALUE = "__default_vm__"
_ALL_VMS_VALUE = "__all_vms__"


def _playbook_for(agent: AgentConfig) -> Path:
    candidate = PLAYBOOKS_DIR / get_agent_class(agent.type).playbook_name()
    if not candidate.exists():
        raise ConfigError(tr("install_agents.script_missing", agent_type=agent.type, path=candidate))
    return candidate


def _format_command(command: List[str]) -> str:
    return " ".join(shlex.quote(part) for part in command)


def _log_failed_command(
    command: List[str],
    result: subprocess.CompletedProcess[str],
    description: str,
) -> None:
    click.echo(tr("install_agents.command_failed", description=description, code=result.returncode), err=True)
    click.echo(tr("install_agents.command_label", command=_format_command(command)), err=True)
    stdout = result.stdout.strip() if result.stdout else ""
    stderr = result.stderr.strip() if result.stderr else ""
    if stdout:
        click.echo(tr("install_agents.stdout_label", output=stdout), err=True)
    if stderr:
        click.echo(tr("install_agents.stderr_label", output=stderr), err=True)


def _run_install_playbook(
    vm: VmConfig,
    playbook_path: Path,
    proxychains: Optional[str] = None,
    *,
    progress: Optional[ProgressManager] = None,
    label: Optional[str] = None,
) -> None:
    ensure_multipass_available()
    try:
        ensure_multipass_collection()
    except AnsibleCollectionError as exc:
        raise MultipassError(str(exc))
    effective_proxychains = resolve_proxychains(vm, proxychains)
    install_command = [
        *ansible_playbook_command(),
        "-i",
        "localhost,",
        "-e",
        f"vm_name={vm.name}",
    ]
    if effective_proxychains:
        install_command.extend(["-e", f"proxychains_url={effective_proxychains}"])
    install_command.append(str(playbook_path))
    if progress and label:
        total = count_playbook_tasks(playbook_path)
        task_id = progress.add_task(label, total=max(total, 1))

        def _handle_progress(current: int, total_tasks: int, task_name: str) -> None:
            effective_total = total_tasks if total_tasks > 0 else total
            progress.update(task_id, description=task_name, completed=min(current, effective_total))

        try:
            result = run_ansible_playbook(
                install_command,
                playbook_path=playbook_path,
                progress_handler=_handle_progress,
                progress_output=progress.print,
            )
        finally:
            progress.remove_task(task_id)
    else:
        result = run_ansible_playbook(
            install_command,
            playbook_path=playbook_path,
        )
    if result.returncode != 0:
        _log_failed_command(install_command, result, tr("install_agents.installer_execution_label"))
        raise MultipassError(tr("install_agents.install_failed", vm_name=vm.name, code=result.returncode))


def _default_vms(agent: AgentConfig, available: Iterable[str]) -> List[str]:
    target_vms = configured_agent_vms(agent, available)
    if not target_vms:
        raise ConfigError(tr("install_agents.no_vms_available"))
    return target_vms


def _select_agent_interactively(agent_names: List[str]) -> Tuple[bool, Optional[str]]:
    choices: list[questionary.QuestionChoice] = [questionary.Choice(tr("interactive.agents_all"), value=_ALL_AGENTS_VALUE)]
    choices.extend(questionary.Choice(name, value=name) for name in agent_names)
    selected = questionary.select(
        tr("interactive.agent_install_select"),
        choices=choices,
        use_shortcuts=True,
    ).ask()
    if selected is None:
        raise click.Abort()
    if selected == _ALL_AGENTS_VALUE:
        return True, None
    return False, str(selected)


def _select_vm_interactively(vms_config: dict[str, VmConfig]) -> Tuple[bool, Optional[str]]:
    default_vm = next(iter(vms_config.keys())) if vms_config else None
    default_label = tr("interactive.vm_default_label")
    if default_vm:
        default_label += f" ({default_vm})"

    choices: list[questionary.QuestionChoice] = [questionary.Choice(default_label, value=_DEFAULT_VM_VALUE)]
    choices.extend(questionary.Choice(name, value=name) for name in vms_config)
    choices.append(questionary.Choice(tr("interactive.vms_all"), value=_ALL_VMS_VALUE))

    selected = questionary.select(
        tr("interactive.agent_install_target"),
        choices=choices,
        use_shortcuts=True,
    ).ask()
    if selected is None:
        raise click.Abort()
    if selected == _ALL_VMS_VALUE:
        return True, None
    if selected == _DEFAULT_VM_VALUE:
        return False, None
    return False, str(selected)


def run_install_agents(
    *,
    agent_name: Optional[str],
    vm: Optional[str],
    all_vms: bool,
    all_agents: bool,
    config_path: Optional[str],
    proxychains: Optional[str],
    debug: bool,
    interactive: bool,
    progress: Optional[ProgressManager] = None,
) -> None:
    with debug_scope(debug):
        if not progress:
            click.echo(tr("install_agents.preparing"))

        if all_agents and agent_name:
            raise click.ClickException(tr("install_agents.agent_conflict"))

        resolved_path = resolve_config_path(Path(config_path) if config_path else None)
        try:
            config = load_config(resolved_path)
            agents_config = load_agents_config(config)
            vms_config = load_vms_config(config)
        except ConfigError as exc:
            raise click.ClickException(str(exc))

        if not agents_config:
            raise click.ClickException(tr("install_agents.no_agents"))

        run_without_args = not agent_name and not vm and not all_agents and not all_vms
        if run_without_args and interactive:
            all_agents, agent_name = _select_agent_interactively(list(agents_config.keys()))
            all_vms, vm = _select_vm_interactively(vms_config)

        agent_names: List[str]
        if all_agents:
            agent_names = list(agents_config.keys())
        else:
            if agent_name:
                agent_names = [agent_name]
            elif len(agents_config) == 1:
                agent_names = [next(iter(agents_config.keys()))]
                click.echo(tr("install_agents.default_agent", agent_name=agent_names[0]))
            else:
                raise click.ClickException(tr("install_agents.agent_required"))

        selected_vms = list(vms_config.keys())
        if vm:
            if vm not in vms_config:
                raise click.ClickException(tr("install_agents.vm_missing", vm_name=vm))
            selected_vms = [vm]

        targets: List[Tuple[str, VmConfig]] = []
        for name in agent_names:
            agent = find_agent(agents_config, name)
            if all_vms:
                for vm_name in vms_config:
                    targets.append((agent.name, vms_config[vm_name]))
            else:
                if vm:
                    chosen_vm = selected_vms[0]
                    if chosen_vm not in vms_config:
                        raise click.ClickException(tr("install_agents.vm_missing", vm_name=chosen_vm))
                    targets.append((agent.name, vms_config[chosen_vm]))
                else:
                    for target_vm_name in _default_vms(agent, vms_config.keys()):
                        if target_vm_name not in vms_config:
                            raise click.ClickException(tr("install_agents.vm_missing", vm_name=target_vm_name))
                        targets.append((agent.name, vms_config[target_vm_name]))

        if progress is None:
            with ProgressManager(debug=debug) as owned_progress:
                _run_install_targets(
                    targets=targets,
                    agents_config=agents_config,
                    proxychains=proxychains,
                    debug=debug,
                    progress=owned_progress,
                    show_overall_task=True,
                )
            return

        _run_install_targets(
            targets=targets,
            agents_config=agents_config,
            proxychains=proxychains,
            debug=debug,
            progress=progress,
            show_overall_task=False,
        )


def _run_install_targets(
    *,
    targets: List[Tuple[str, VmConfig]],
    agents_config: dict[str, AgentConfig],
    proxychains: Optional[str],
    debug: bool,
    progress: ProgressManager,
    show_overall_task: bool,
) -> None:
    overall_task = progress.add_task(tr("progress.install_agents_title"), total=len(targets)) if show_overall_task else None
    for target_agent_name, target_vm in targets:
        agent = find_agent(agents_config, target_agent_name)
        playbook_path = _playbook_for(agent)
        if debug:
            click.echo(
                tr(
                    "install_agents.installing",
                    agent_name=agent.name,
                    agent_type=agent.type,
                    vm_name=target_vm.name,
                    script=playbook_path.name,
                )
            )
        try:
            proxychains_override = proxychains
            if proxychains_override is None and agent.proxychains_defined:
                proxychains_override = agent.proxychains if agent.proxychains is not None else ""
            install_label = tr(
                "progress.install_agent_target",
                agent_name=agent.name,
                agent_type=agent.type,
                vm_name=target_vm.name,
            )
            if overall_task is not None:
                progress.update(overall_task, description=install_label)
            _run_install_playbook(
                target_vm,
                playbook_path,
                proxychains=proxychains_override,
                progress=progress,
                label=install_label,
            )
        except (MultipassError, ConfigError) as exc:
            raise click.ClickException(str(exc))
        if debug:
            click.echo(
                tr(
                    "install_agents.installed",
                    agent_name=agent.name,
                    agent_type=agent.type,
                        vm_name=target_vm.name,
                    )
                )
        if overall_task is not None:
            progress.advance(overall_task)
    if overall_task is not None:
        progress.remove_task(overall_task)


@click.command(name="install-agents", help=tr("install_agents.command_help"))
@non_interactive_option
@click.argument("agent_name", required=False)
@click.argument("vm", required=False)
@click.option("--all-vms", is_flag=True, help=tr("install_agents.option_all_vms"))
@click.option("--all-agents", is_flag=True, help=tr("install_agents.option_all_agents"))
@click.option(
    "config_path",
    "--config",
    type=click.Path(dir_okay=False, exists=False, path_type=str),
    envvar="CONFIG_PATH",
    default=None,
    help=tr("config.option_path"),
)
@click.option(
    "--proxychains",
    default=None,
    show_default=False,
    help=tr("install_agents.option_proxychains"),
)
@debug_option
def install_agents_command(
    agent_name: Optional[str],
    vm: Optional[str],
    all_vms: bool,
    all_agents: bool,
    config_path: Optional[str],
    proxychains: Optional[str],
    debug: bool,
    non_interactive: bool,
) -> None:
    """Install configured agents into Multipass VMs."""
    interactive = is_interactive_terminal() and not non_interactive
    run_install_agents(
        agent_name=agent_name,
        vm=vm,
        all_vms=all_vms,
        all_agents=all_agents,
        config_path=config_path,
        proxychains=proxychains,
        debug=debug,
        interactive=interactive,
    )
