"""Vagrant VM lifecycle and WinRM tools."""

import asyncio
import json
import os
import shlex
import subprocess
import tempfile
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Annotated

import libtmux
import vagrant
import winrm
from fastmcp import Context
from fastmcp.dependencies import CurrentContext, Progress
from fastmcp.exceptions import ToolError
from fastmcp.tools import tool
from pydantic import Field

from mcp_vm_blackbox._helpers import _resolve, _ssh_run, default_blackbox_data_root, effective_root
from mcp_vm_blackbox.models import VagrantCommandResult, VagrantDetachedResult
from mcp_vm_blackbox.targets import TargetConfig

_DEFAULT_WINRM_PASSWORD = "vagrant"  # noqa: S105
_KEY_VALUE_SPLIT_MAX = 2


def _parse_winrm_config(output: str) -> dict[str, str]:
    """Parse the output of ``vagrant winrm-config`` into a config dict."""
    conf: dict[str, str] = {}
    started = False
    for line in output.splitlines():
        if line.strip().startswith("Host ") and not started:
            started = True
        if not started or not line.strip() or line.strip().startswith("#"):
            continue
        parts = line.strip().split(None, 1)
        if len(parts) == _KEY_VALUE_SPLIT_MAX:
            key, value = parts
            conf[key] = value.strip('"')
    return conf


def _get_winrm_config(v: vagrant.Vagrant, vm_name: str) -> dict[str, str]:
    """Get WinRM connection config for a VM via python-vagrant."""
    try:
        output = v._run_vagrant_command(["winrm-config", vm_name])  # noqa: SLF001
    except Exception as exc:
        raise ToolError(
            f"Failed to get WinRM config for VM {vm_name!r}. Ensure the VM is up and uses the WinRM communicator: {exc}"
        ) from exc
    return _parse_winrm_config(output)


def _winrm_run_sync(host: str, port: int, username: str, password: str, command: str) -> tuple[str, str, int]:
    """Execute a PowerShell command via pywinrm (synchronous)."""
    session = winrm.Session(f"http://{host}:{port}/wsman", auth=(username, password), transport="ntlm")
    result = session.run_ps(command)
    stdout = (result.std_out or b"").decode("utf-8", errors="replace")
    stderr = (result.std_err or b"").decode("utf-8", errors="replace")
    return stdout, stderr, result.status_code or 0


async def _winrm_run_local(
    root: str, vm_name: str, command: str, extra_env: dict[str, str] | None, timeout_s: int, start: float
) -> tuple[str, str, int, float]:
    """Run a command on a Windows Vagrant VM via pywinrm (no subprocess)."""
    env = {**os.environ, **(extra_env or {})}
    v = vagrant.Vagrant(root=root, quiet_stdout=True, quiet_stderr=True, env=env)
    conf = _get_winrm_config(v, vm_name)
    host = conf.get("HostName", "127.0.0.1")
    port = int(conf.get("Port", "5985"))
    username = conf.get("User", "vagrant")
    password = conf.get("Password", _DEFAULT_WINRM_PASSWORD)
    try:
        stdout, stderr, exit_code = await asyncio.wait_for(
            asyncio.to_thread(_winrm_run_sync, host, port, username, password, command), timeout=timeout_s
        )
        return stdout, stderr, exit_code, round(time.monotonic() - start, 1)
    except TimeoutError:
        return ("", f"timed out after {timeout_s}s", -1, round(time.monotonic() - start, 1))


def _extract_provision_with(args: list[str]) -> list[str] | None:
    """Extract a ``--provision-with`` value from *args*."""
    if "--provision-with" not in args:
        return None
    idx = args.index("--provision-with")
    if idx + 1 < len(args):
        return [args[idx + 1]]
    return None


def _vagrant_dispatch_local(v: vagrant.Vagrant, args: list[str]) -> str | list[object] | None:
    """Dispatch a vagrant sub-command to the python-vagrant library."""
    sub_cmd = args[0] if args else ""
    vm_name: str | None = args[1] if len(args) > 1 else None

    match sub_cmd:
        case "status":
            return v.status(vm_name=vm_name)
        case "up":
            return v.up(vm_name=vm_name, provision_with=_extract_provision_with(args))
        case "provision":
            return v.provision(vm_name=vm_name, provision_with=_extract_provision_with(args))
        case "destroy":
            return v.destroy(vm_name=vm_name)
        case "halt":
            return v.halt(vm_name=vm_name)
        case _:
            raise ToolError(
                f"Unsupported vagrant subcommand: {sub_cmd!r}. "
                f"Supported via python-vagrant: up, halt, status, provision, destroy. "
                f"Add explicit handling for {sub_cmd!r} if needed."
            )


def _vagrant_run_local_sync(
    root: str, args: list[str], extra_env: dict[str, str] | None, start: float
) -> tuple[str, str, int, float]:
    """Run a vagrant command locally, capturing stdout and stderr via subprocess.run.

    Uses python-vagrant only to resolve the vagrant executable path and build the
    command list. The subprocess is run directly with capture_output=True so that
    output is always captured regardless of TTY buffering — python-vagrant's
    out_cm/err_cm file-handle approach suffers from full-buffering when stdout is
    not a TTY, which causes output to be lost when vagrant exits non-zero.
    """
    env = {**os.environ, **(extra_env or {})}

    sub_cmd = args[0] if args else ""

    # status uses _run_vagrant_command (returns stdout via check_output) —
    # delegate to python-vagrant for its structured return value.
    if sub_cmd == "status":
        v = vagrant.Vagrant(root=root, quiet_stdout=True, quiet_stderr=True, env=env)
        try:
            dispatch_result = _vagrant_dispatch_local(v, args)
            stdout = str(dispatch_result) if dispatch_result is not None else ""
            return stdout, "", 0, round(time.monotonic() - start, 1)
        except ToolError:
            raise
        except Exception as exc:
            raise ToolError(f"vagrant 'status' failed: {exc}") from exc

    # For up, provision, destroy, halt — run via subprocess.run with
    # capture_output=True.  This guarantees stdout and stderr are always
    # available on non-zero exit, regardless of TTY buffering.
    v = vagrant.Vagrant(root=root, env=env)
    command = v._make_vagrant_command(args)  # noqa: SLF001
    result = subprocess.run(command, check=False, cwd=root, env=env, capture_output=True)
    stdout = result.stdout.decode("utf-8", errors="replace")
    stderr = result.stderr.decode("utf-8", errors="replace")
    if result.returncode != 0:
        parts = [f"vagrant {sub_cmd!r} failed (exit {result.returncode})"]
        if stdout.strip():
            parts.append(f"stdout:\n{stdout.rstrip()}")
        if stderr.strip():
            parts.append(f"stderr:\n{stderr.rstrip()}")
        raise ToolError("\n".join(parts))
    return stdout, stderr, result.returncode, round(time.monotonic() - start, 1)


async def _vagrant_run_remote(
    cfg: TargetConfig, args: list[str], root: str, extra_env: dict[str, str] | None, timeout_s: int, start: float
) -> tuple[str, str, int, float]:
    """Forward a vagrant command to a remote host via asyncssh."""
    env_prefix = " ".join(f"{k}={shlex.quote(v)}" for k, v in (extra_env or {}).items())
    vagrant_args = " ".join(shlex.quote(a) for a in args)
    cmd_str = f"cd {shlex.quote(root)} && {env_prefix} vagrant {vagrant_args}".strip()
    try:
        result = await asyncio.wait_for(_ssh_run(cfg, ["bash", "-c", cmd_str]), timeout=timeout_s)
    except TimeoutError:
        return ("", f"timed out after {timeout_s}s", -1, round(time.monotonic() - start, 1))
    return (result.stdout, result.stderr, result.returncode, round(time.monotonic() - start, 1))


async def _vagrant_run(
    cfg: TargetConfig,
    args: list[str],
    timeout_s: int = 600,
    extra_env: dict[str, str] | None = None,
    repo_root: str | None = None,
) -> tuple[str, str, int, float]:
    """Run a vagrant command, returning (stdout, stderr, exit_code, duration_s)."""
    root = effective_root(repo_root)
    start = time.monotonic()

    if cfg.is_remote:
        return await _vagrant_run_remote(cfg, args, root, extra_env, timeout_s, start)

    sub_cmd = args[0] if args else ""
    if sub_cmd == "winrm":
        vm_name = args[1] if len(args) > 1 else "default"
        cmd = ""
        if "-c" in args:
            idx = args.index("-c")
            if idx + 1 < len(args):
                cmd = args[idx + 1]
        return await _winrm_run_local(root, vm_name, cmd, extra_env, timeout_s, start)

    return await asyncio.to_thread(_vagrant_run_local_sync, root, args, extra_env, start)


async def _vagrant_run_detached(
    cfg: TargetConfig, args: list[str], vm: str, extra_env: dict[str, str] | None = None, repo_root: str | None = None
) -> VagrantDetachedResult:
    """Launch a vagrant command in a background session and return immediately."""
    timestamp = datetime.now(tz=UTC).strftime("%Y%m%d-%H%M%S")
    safe_vm = vm.replace("/", "-").replace(" ", "-")
    session_name = f"vagrant-{safe_vm}-{timestamp}"
    log_path = f"{tempfile.gettempdir()}/{session_name}.log"

    root = effective_root(repo_root)
    env_prefix = " ".join(f"{k}={shlex.quote(v)}" for k, v in (extra_env or {}).items())
    vagrant_args = " ".join(shlex.quote(a) for a in args)
    vagrant_cmd = (
        f"cd {shlex.quote(root)} && {env_prefix} vagrant {vagrant_args} 2>&1 | tee {shlex.quote(log_path)}".strip()
    )

    if cfg.ssh_host:
        launch_inner = f"tmux new-session -d -s {shlex.quote(session_name)} {shlex.quote(vagrant_cmd)}"
        result = await _ssh_run(cfg, ["bash", "-c", launch_inner])
        if result.returncode != 0:
            raise ToolError(f"Failed to launch detached vagrant session on {cfg.ssh_host}: {result.stderr.strip()}")
    else:
        try:
            server = libtmux.Server()
            server.new_session(session_name=session_name, window_command=vagrant_cmd, attach=False)
        except Exception as exc:
            raise ToolError(
                f"tmux is required for detached vagrant operations but could not be started: {exc}. "
                f"Install tmux or run the command without detached=True."
            ) from exc

    return VagrantDetachedResult(
        detached=True,
        session_name=session_name,
        log_path=log_path,
        tail_command=f"tail -f {log_path}",
        attach_command=f"tmux attach -t {session_name}",
        kill_command=f"tmux kill-session -t {session_name}",
    )


def _write_run_status(composite_key: str, status: str, root: Path) -> None:
    """Persist run status to .mcp-vm-blackbox/run-status/{composite_key}.json."""
    parts = composite_key.split("/", 1)
    vm_name = parts[0]
    scenario = parts[1] if len(parts) > 1 else "default"
    path = root / "run-status" / vm_name / f"{scenario}.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=path.parent, suffix=".tmp")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            json.dump({"status": status}, f)
        Path(tmp).rename(path)
    except Exception:
        Path(tmp).unlink(missing_ok=True)
        raise


async def _mark_run_complete(composite_key: str, ctx: Context, store_root: Path | None = None) -> None:
    """Mark a scenario run as completed and open the verification gate."""
    root = store_root if store_root is not None else default_blackbox_data_root()
    _write_run_status(composite_key, "completed", root)

    try:
        await ctx.set_state(f"run_status:{composite_key}", "completed")
        await ctx.enable_components(tags={"verify"})
    except RuntimeError as exc:
        if "session_id is not available" not in str(exc):
            raise


@tool
async def vagrant_status(
    target: Annotated[
        str, Field(description='Named target ("local", "ci") or raw SSH string. Defaults to "local".')
    ] = "local",
    repo_root: Annotated[
        str | None,
        Field(
            description="Absolute path to the repository root containing the Vagrantfile. "
            "When None, uses the default repo root detected from this module's location."
        ),
    ] = None,
    timeout_s: Annotated[int, Field(description="Timeout in seconds.", ge=1)] = 30,
) -> VagrantCommandResult:
    """Show the status of all Vagrant VMs defined in the Vagrantfile."""
    root = effective_root(repo_root)
    cfg = _resolve(target)
    stdout, stderr, exit_code, duration_s = await _vagrant_run(cfg, ["status"], timeout_s=timeout_s, repo_root=root)
    if exit_code == -1:
        return VagrantCommandResult(exit_code=exit_code, duration_s=duration_s, error=stderr)
    return VagrantCommandResult(stdout=stdout, stderr=stderr, exit_code=exit_code, duration_s=duration_s)


@tool(task=True)
async def vagrant_up(
    vm: Annotated[
        str, Field(description='Vagrant VM name to start, or "all" to start every VM in the Vagrantfile.')
    ] = "all",
    scenario_name: Annotated[
        str, Field(description="Scenario name used to key the goal-state verification gate. Defaults to 'default'.")
    ] = "default",
    target: Annotated[str, Field(description='Named target or raw SSH string. Defaults to "local".')] = "local",
    repo_root: Annotated[
        str | None, Field(description="Absolute path to the repository root. When None, uses the default.")
    ] = None,
    timeout_s: Annotated[int, Field(description="Timeout in seconds. Defaults to 600.", ge=1)] = 600,
    progress: Progress = Progress(),  # noqa: B008
    ctx: Context = CurrentContext(),  # noqa: B008
) -> VagrantCommandResult:
    """Start one or all Vagrant VMs."""
    await progress.set_message(f"Starting vagrant up{'' if vm == 'all' else f' {vm}'}...")
    root = effective_root(repo_root)
    cfg = _resolve(target)
    args = ["up"] if vm == "all" else ["up", vm]
    stdout, stderr, exit_code, duration_s = await _vagrant_run(cfg, args, timeout_s=timeout_s, repo_root=root)
    await progress.set_message("vagrant up complete")
    if exit_code == -1:
        return VagrantCommandResult(exit_code=exit_code, duration_s=duration_s, error=stderr)
    await _mark_run_complete(f"{vm}/{scenario_name}", ctx)
    return VagrantCommandResult(stdout=stdout, stderr=stderr, exit_code=exit_code, duration_s=duration_s)


@tool(task=True)
async def vagrant_provision(
    vm: Annotated[str, Field(description="Vagrant VM name to provision (required).")],
    scenario_name: Annotated[
        str, Field(description="Scenario name used to key the goal-state verification gate. Defaults to 'default'.")
    ] = "default",
    provisioner: Annotated[
        str | None,
        Field(description='Optional provisioner name (e.g. "ansible_local"). When None, all provisioners run.'),
    ] = None,
    target: Annotated[str, Field(description='Named target or raw SSH string. Defaults to "local".')] = "local",
    repo_root: Annotated[
        str | None, Field(description="Absolute path to the repository root. When None, uses the default.")
    ] = None,
    timeout_s: Annotated[int, Field(description="Timeout in seconds. Defaults to 1800.", ge=1)] = 1800,
    progress: Progress = Progress(),  # noqa: B008
    ctx: Context = CurrentContext(),  # noqa: B008
) -> VagrantCommandResult:
    """Run Vagrant provisioners on a specific VM."""
    await progress.set_message(f"Starting vagrant provision {vm}...")
    root = effective_root(repo_root)
    cfg = _resolve(target)
    args = ["provision", vm]
    if provisioner:
        args += ["--provision-with", provisioner]
    stdout, stderr, exit_code, duration_s = await _vagrant_run(cfg, args, timeout_s=timeout_s, repo_root=root)
    await progress.set_message("vagrant provision complete")
    if exit_code == -1:
        return VagrantCommandResult(exit_code=exit_code, duration_s=duration_s, error=stderr)
    await _mark_run_complete(f"{vm}/{scenario_name}", ctx)
    return VagrantCommandResult(stdout=stdout, stderr=stderr, exit_code=exit_code, duration_s=duration_s)


@tool
async def vagrant_destroy(
    vm: Annotated[str, Field(description="Vagrant VM name to destroy (required).")],
    force: Annotated[bool, Field(description='When True (default), passes "-f" to skip confirmation.')] = True,
    target: Annotated[str, Field(description='Named target or raw SSH string. Defaults to "local".')] = "local",
    repo_root: Annotated[
        str | None, Field(description="Absolute path to the repository root. When None, uses the default.")
    ] = None,
    timeout_s: Annotated[int, Field(description="Timeout in seconds. Defaults to 600.", ge=1)] = 600,
    detached: Annotated[
        bool, Field(description="When True, run in a background tmux session and return immediately.")
    ] = False,
) -> VagrantCommandResult | VagrantDetachedResult:
    """Destroy a Vagrant VM, stopping and deleting all associated resources."""
    from mcp_vm_blackbox.job_store import default_job_store  # noqa: PLC0415

    all_jobs = await default_job_store.list_by_vm(vm)
    active_jobs = [j for j in all_jobs if j.status == "active"]
    if active_jobs:
        blocker = active_jobs[0]
        raise ToolError(
            f"Cannot destroy VM '{vm}': active job '{blocker.job_id}' "
            f"(scenario={blocker.scenario_name}, last_agent={blocker.last_agent_id}) "
            f"is still running. Complete or pause the job first."
        )

    root = effective_root(repo_root)
    cfg = _resolve(target)
    args = ["destroy", vm]
    if force:
        args.append("-f")

    if detached:
        return await _vagrant_run_detached(cfg, args, vm=vm, repo_root=root)

    stdout, stderr, exit_code, duration_s = await _vagrant_run(cfg, args, timeout_s=timeout_s, repo_root=root)
    if exit_code == -1:
        return VagrantCommandResult(exit_code=exit_code, duration_s=duration_s, error=stderr)
    return VagrantCommandResult(stdout=stdout, stderr=stderr, exit_code=exit_code, duration_s=duration_s)


@tool
async def vagrant_winrm(
    vm: Annotated[str, Field(description='Vagrant VM name (e.g. "my-vm").')],
    command: Annotated[str, Field(description="PowerShell or cmd command to execute on the VM.")],
    target: Annotated[
        str, Field(description='Accepted for API consistency. Only "local" is meaningful here.')
    ] = "local",
    repo_root: Annotated[
        str | None,
        Field(
            description="Absolute path to the repository root containing the Vagrantfile. When None, uses the default."
        ),
    ] = None,
    timeout_s: Annotated[int, Field(description="Timeout in seconds. Defaults to 30.", ge=1)] = 30,
) -> VagrantCommandResult:
    """Run a command on a Windows Vagrant VM via WinRM."""
    cfg = _resolve(target)
    args = ["winrm", vm, "-c", command]
    stdout, stderr, exit_code, duration_s = await _vagrant_run(
        cfg, args, timeout_s=timeout_s, repo_root=effective_root(repo_root)
    )
    if exit_code == -1:
        return VagrantCommandResult(exit_code=exit_code, duration_s=duration_s, error=stderr)
    return VagrantCommandResult(stdout=stdout, stderr=stderr, exit_code=exit_code, duration_s=duration_s)
