"""Provisioning engine — installs and configures agentsh inside a sandbox.

See spec Section 8 (8.1-8.4) for the full provisioning sequence.
"""

from __future__ import annotations

import asyncio
import json
import shlex
from dataclasses import dataclass
from typing import TYPE_CHECKING

from .types import (
    ExecOpts,
    SecurityMode,
)
from .integrity import get_checksum, build_verify_command, binary_url, PINNED_VERSION
from .config import generate_server_config
from .errors import IntegrityError, ProvisioningError
from ..policies.serialize import serialize_policy, system_policy_yaml
from ..policies.presets import agent_default
from ..policies.schema import validate_policy

if TYPE_CHECKING:
    from .types import SandboxAdapter, SecureConfig

SECURITY_MODE_RANK: dict[SecurityMode, int] = {
    "full": 5,
    "ptrace": 4,
    "landlock": 3,
    "landlock-only": 2,
    "minimal": 1,
}

AGENTSH_PATHS = ["/usr/local/bin/agentsh", "/usr/bin/agentsh"]

_HEALTH_CHECK_ATTEMPTS = 10
_HEALTH_CHECK_INTERVAL = 0.5  # seconds


@dataclass(frozen=True, slots=True)
class ProvisionResult:
    session_id: str
    security_mode: SecurityMode
    passthrough: bool = False


# ─── Main entry point ──────────────────────────────────


async def provision(
    adapter: SandboxAdapter,
    config: SecureConfig | None = None,
) -> ProvisionResult:
    """Orchestrate the full provisioning sequence.

    For ``install_strategy="running"``, returns immediately with passthrough mode.
    For ``download`` or ``preinstalled``, runs Phases 1-4 from the spec.
    """
    from .types import SecureConfig as _SC

    cfg = config or _SC()
    policy = validate_policy(cfg.policy) if cfg.policy else agent_default()

    # ── Fast path: running strategy ──
    if cfg.install_strategy == "running":
        # agentsh is already fully provisioned and running with the shell shim.
        # Skip install, shim, policy, config, server startup, and security detection.
        # Default securityMode to 'full' (override via securityMode config).
        security_mode: SecurityMode = cfg.security_mode or "full"

        # Check minimum security mode even in running strategy
        if cfg.minimum_security_mode:
            _check_minimum_security_mode(security_mode, cfg.minimum_security_mode)

        # Health check — verify the server is reachable
        await _health_check(adapter)

        # Read session ID from config or from env
        session_id = cfg.session_id
        if not session_id:
            env_result = await adapter.exec(
                "bash", ["-c", "echo $AGENTSH_SESSION_ID"]
            )
            session_id = env_result.stdout.strip()

        if not session_id:
            raise ProvisioningError(
                phase="running_session",
                command="echo $AGENTSH_SESSION_ID",
                stderr="AGENTSH_SESSION_ID not set — running strategy requires a pre-created session",
            )

        return ProvisionResult(
            session_id=session_id,
            security_mode=security_mode,
            passthrough=True,
        )

    # ── Phase 1: Binary Installation ──

    # 1a. Detect architecture
    arch = cfg.agentsh_arch or await _detect_arch(adapter)

    # 1b. Resolve version
    version = cfg.agentsh_version or PINNED_VERSION

    # 1c. Install binary (download or verify preinstalled)
    if cfg.install_strategy == "download":
        binary_path = AGENTSH_PATHS[0]  # install to /usr/local/bin/agentsh
        await _download_tarball(adapter, version, arch, cfg)
        if not cfg.skip_integrity_check:
            await _verify_integrity(adapter, version, arch, cfg)
        await _extract_and_install(adapter)
    elif cfg.install_strategy == "upload":
        binary_path = AGENTSH_PATHS[0]
        await _upload_tarball(adapter, version, arch, cfg)
        if not cfg.skip_integrity_check:
            await _verify_integrity(adapter, version, arch, cfg)
        await _extract_and_install(adapter)
    else:
        # preinstalled: probe known paths to find the binary
        binary_path = await _find_preinstalled_binary(adapter)

    # 1d. Detect security mode
    security_mode = await _detect_security_mode(adapter, binary_path)

    # 1f. Check minimum security mode
    if cfg.minimum_security_mode:
        _check_minimum_security_mode(security_mode, cfg.minimum_security_mode)

    # 1g. Install shell shim (skip when ptrace handles enforcement)
    if not cfg.skip_shim:
        await _install_shell_shim(adapter, binary_path)

    # ── Phase 2: Policy & Config ──

    # 2a. Create config directories
    await _create_config_dirs(adapter)

    # Temporarily make config dir writable for writing files
    await adapter.exec("chmod", ["-R", "777", "/etc/agentsh/"], opts=ExecOpts(sudo=True))

    # 2b. Write system policy (self-protection rules)
    sys_policy = system_policy_yaml()
    await adapter.write_file(
        "/etc/agentsh/system/policy.yml",
        sys_policy,
        sudo=True,
    )

    # 2c. Write user policy
    user_policy = serialize_policy(policy)
    await adapter.write_file(
        "/etc/agentsh/policy.yml",
        user_policy,
        sudo=True,
    )

    # 2d. Write server config
    server_config = generate_server_config(cfg)
    await adapter.write_file(
        "/etc/agentsh/config.yml",
        server_config,
        sudo=True,
    )

    # 2e. Set permissions
    await _set_permissions(adapter)

    # 2f. Ensure workspace directory exists
    await adapter.exec("mkdir", ["-p", cfg.workspace], opts=ExecOpts(sudo=True))

    # Create sessions dir
    await adapter.exec("mkdir", ["-p", cfg.workspace, "/var/lib/agentsh/sessions"], opts=ExecOpts(sudo=True))
    await adapter.exec("chmod", ["755", "/var/lib/agentsh", "/var/lib/agentsh/sessions"], opts=ExecOpts(sudo=True))

    # Allow non-root users to access FUSE mounts
    await adapter.exec("sh", [
        "-c",
        'grep -q user_allow_other /etc/fuse.conf 2>/dev/null || echo "user_allow_other" >> /etc/fuse.conf',
    ], opts=ExecOpts(sudo=True))

    # ── Phase 3: Server Startup ──

    # 3a. Start agentsh server (detached)
    await _start_server(adapter, binary_path)

    # 3b. Health check
    await _health_check(adapter)

    # 3c. Create session
    session_id = await _create_session(adapter, binary_path, cfg.workspace)

    # Make session dir readable by non-root users
    await adapter.exec("chmod", ["-R", "755", "/var/lib/agentsh/sessions/"], opts=ExecOpts(sudo=True))

    # Set trace context if provided
    from .trace import resolve_trace_parent
    effective_trace = resolve_trace_parent(cfg.trace_parent)
    if effective_trace:
        await adapter.exec("curl", [
            "-X", "PUT",
            f"http://127.0.0.1:18080/sessions/{session_id}/trace-context",
            "-H", "Content-Type: application/json",
            "-d", json.dumps({"traceparent": effective_trace}),
        ])

    # ── Phase 4: Handoff ──
    return ProvisionResult(
        session_id=session_id,
        security_mode=security_mode,
    )


# ─── Helper functions ──────────────────────────────────


async def _detect_arch(adapter: SandboxAdapter) -> str:
    """Detect CPU architecture via ``uname -m`` and map to agentsh arch string."""
    result = await adapter.exec("uname", ["-m"])
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="detect_arch",
            command="uname -m",
            stderr=result.stderr,
        )
    raw = result.stdout.strip()
    mapping = {
        "x86_64": "linux_amd64",
        "aarch64": "linux_arm64",
        "arm64": "linux_arm64",
    }
    arch = mapping.get(raw)
    if not arch:
        raise ProvisioningError(
            phase="detect_arch",
            command="uname -m",
            stderr=f"Unsupported architecture: {raw}",
        )
    return arch


async def _find_preinstalled_binary(adapter: SandboxAdapter) -> str:
    """Probe AGENTSH_PATHS to find a preinstalled agentsh binary."""
    for path in AGENTSH_PATHS:
        if await adapter.file_exists(path):
            return path
    raise ProvisioningError(
        phase="find_binary",
        command=f"file_exists({AGENTSH_PATHS})",
        stderr=f"agentsh binary not found at any of: {', '.join(AGENTSH_PATHS)}",
    )


async def _download_tarball(
    adapter: SandboxAdapter,
    version: str,
    arch: str,
    cfg: SecureConfig,
) -> None:
    """Download agentsh tarball inside sandbox to /tmp/agentsh.tar.gz."""
    url = binary_url(version, arch, cfg.agentsh_binary_url)
    safe_url = shlex.quote(url)
    download_cmd = f"curl -fsSL {safe_url} -o /tmp/agentsh.tar.gz"
    result = await adapter.exec(
        "bash",
        ["-c", download_cmd],
        opts=ExecOpts(sudo=True),
    )
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="download_binary",
            command=download_cmd,
            stderr=result.stderr,
        )


async def _upload_tarball(
    adapter: SandboxAdapter,
    version: str,
    arch: str,
    cfg: SecureConfig,
) -> None:
    """Download agentsh on host via httpx, upload to /tmp/agentsh.tar.gz."""
    import httpx

    url = binary_url(version, arch, cfg.agentsh_binary_url)
    async with httpx.AsyncClient(follow_redirects=True) as client:
        resp = await client.get(url)
        resp.raise_for_status()

    # Write tarball to sandbox
    await adapter.write_file("/tmp/agentsh.tar.gz", resp.content, sudo=True)


async def _extract_and_install(adapter: SandboxAdapter) -> None:
    """Extract tarball and install agentsh binaries."""
    install_cmd = (
        "tar -xzf /tmp/agentsh.tar.gz -C /tmp/ && "
        "install -m 0755 /tmp/agentsh /usr/local/bin/agentsh && "
        "install -m 0755 /tmp/agentsh-shell-shim /usr/bin/agentsh-shell-shim && "
        "install -m 0755 /tmp/agentsh-unixwrap /usr/local/bin/agentsh-unixwrap && "
        "rm -f /tmp/agentsh.tar.gz"
    )
    result = await adapter.exec("bash", ["-c", install_cmd], opts=ExecOpts(sudo=True))
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="install_binary",
            command=install_cmd,
            stderr=result.stderr,
        )


async def _verify_integrity(
    adapter: SandboxAdapter,
    version: str,
    arch: str,
    cfg: SecureConfig,
) -> None:
    """Verify SHA256 checksum of the downloaded tarball."""
    expected = cfg.agentsh_checksum or get_checksum(version, arch)
    verify_cmd = build_verify_command("/tmp/agentsh.tar.gz", expected)
    result = await adapter.exec("bash", ["-c", verify_cmd])
    if result.exit_code != 0:
        raise IntegrityError(
            expected=expected,
            actual="unknown",
            message="Tarball integrity check failed for /tmp/agentsh.tar.gz",
        )


async def _detect_security_mode(
    adapter: SandboxAdapter,
    binary_path: str,
) -> SecurityMode:
    """Run ``agentsh detect --output json`` to determine security capabilities."""
    result = await adapter.exec(binary_path, ["detect", "--output", "json"])
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="detect_security_mode",
            command=f"{binary_path} detect --output json",
            stderr=result.stderr,
        )
    try:
        # agentsh detect outputs JSON to stderr (not stdout)
        json_output = result.stderr or result.stdout
        data = json.loads(json_output)
        mode: SecurityMode = data["security_mode"]
        if mode not in SECURITY_MODE_RANK:
            raise ValueError(f"Unknown security mode: {mode}")
        return mode
    except (json.JSONDecodeError, KeyError, ValueError) as exc:
        raise ProvisioningError(
            phase="detect_security_mode",
            command=f"{binary_path} detect --output json",
            stderr=f"Failed to parse detect output: {exc}",
        ) from exc


def _check_minimum_security_mode(
    detected: SecurityMode,
    minimum: SecurityMode,
) -> None:
    """Raise ProvisioningError if detected mode is weaker than minimum."""
    if SECURITY_MODE_RANK[detected] < SECURITY_MODE_RANK[minimum]:
        raise ProvisioningError(
            phase="security_mode_check",
            command="agentsh detect",
            stderr=(
                f"Detected security mode '{detected}' (rank {SECURITY_MODE_RANK[detected]}) "
                f"is weaker than required minimum '{minimum}' "
                f"(rank {SECURITY_MODE_RANK[minimum]})"
            ),
        )


async def _install_shell_shim(
    adapter: SandboxAdapter,
    binary_path: str,
) -> None:
    """Install the agentsh shell shim."""
    result = await adapter.exec(
        binary_path,
        [
            "shim", "install-shell",
            "--root", "/",
            "--shim", "/usr/bin/agentsh-shell-shim",
            "--bash",
            "--i-understand-this-modifies-the-host",
        ],
        opts=ExecOpts(sudo=True),
    )
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="install_shell_shim",
            command=f"{binary_path} shim install-shell",
            stderr=result.stderr,
        )


async def _create_config_dirs(adapter: SandboxAdapter) -> None:
    """Create /etc/agentsh/ and /etc/agentsh/system/ directories."""
    result = await adapter.exec(
        "mkdir",
        ["-p", "/etc/agentsh/system"],
        opts=ExecOpts(sudo=True),
    )
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="create_config_dirs",
            command="mkdir -p /etc/agentsh/system",
            stderr=result.stderr,
        )


async def _set_permissions(adapter: SandboxAdapter) -> None:
    """Set restrictive permissions on config files."""
    # Dirs: 555, Files: 444, Owner: root:root
    await adapter.exec(
        "find", ["/etc/agentsh", "-type", "d", "-exec", "chmod", "555", "{}", "+"],
        opts=ExecOpts(sudo=True),
    )
    await adapter.exec(
        "find", ["/etc/agentsh", "-type", "f", "-exec", "chmod", "444", "{}", "+"],
        opts=ExecOpts(sudo=True),
    )
    await adapter.exec(
        "chown", ["-R", "root:root", "/etc/agentsh/"],
        opts=ExecOpts(sudo=True),
    )


async def _start_server(
    adapter: SandboxAdapter,
    binary_path: str,
) -> None:
    """Start the agentsh server in detached mode."""
    result = await adapter.exec(
        binary_path,
        ["server", "--config", "/etc/agentsh/config.yml"],
        opts=ExecOpts(sudo=True, detached=True),
    )
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="start_server",
            command=f"{binary_path} server --config /etc/agentsh/config.yml",
            stderr=result.stderr,
        )


async def _health_check(adapter: SandboxAdapter) -> None:
    """Poll the agentsh server health endpoint until it responds."""
    for attempt in range(_HEALTH_CHECK_ATTEMPTS):
        result = await adapter.exec(
            "curl",
            ["-sf", "http://127.0.0.1:18080/health"],
        )
        if result.exit_code == 0:
            return
        if attempt < _HEALTH_CHECK_ATTEMPTS - 1:
            await asyncio.sleep(_HEALTH_CHECK_INTERVAL)

    raise ProvisioningError(
        phase="health_check",
        command="curl -sf http://127.0.0.1:18080/health",
        stderr=f"Server failed to become healthy after {_HEALTH_CHECK_ATTEMPTS} attempts",
    )


async def _create_session(
    adapter: SandboxAdapter,
    binary_path: str,
    workspace: str,
) -> str:
    """Create an agentsh session and return the session ID."""
    result = await adapter.exec(
        binary_path,
        ["session", "create", "--workspace", workspace, "--policy", "policy"],
    )
    if result.exit_code != 0:
        raise ProvisioningError(
            phase="create_session",
            command=f"{binary_path} session create --workspace {workspace} --policy policy",
            stderr=result.stderr,
        )
    try:
        data = json.loads(result.stdout)
        session_id = data["session_id"]
        if not isinstance(session_id, str) or not session_id:
            raise ValueError("Empty or invalid session_id")
        return session_id
    except (json.JSONDecodeError, KeyError, ValueError):
        # Fallback: parse text output like "Session session-xxx started"
        import re
        match = re.search(r"Session\s+(session-[^\s]+)", result.stdout)
        if match:
            return match[1]
        raise ProvisioningError(
            phase="create_session",
            command=f"{binary_path} session create --workspace {workspace} --policy policy",
            stderr=f"Failed to parse session output: {result.stdout}",
        )
