"""Map a SandboxPolicy to a ContainerConfig for the Docker runtime."""

from __future__ import annotations

import base64
import re
import tomllib
from pathlib import Path

from hatchdx.sandbox.models import ContainerConfig, SandboxPolicy


def policy_to_container_config(
    policy: SandboxPolicy,
    *,
    image: str = "",
    project_path: str = "",
) -> ContainerConfig:
    """Convert a high-level ``SandboxPolicy`` into a ``ContainerConfig``.

    The returned config encodes the policy's network, filesystem, and resource
    constraints in Docker-native terms (network mode, volume mounts, tmpfs,
    CPU/memory limits, etc.).

    Parameters
    ----------
    policy:
        The sandbox security policy to translate.
    image:
        Docker image name/tag to use.
    project_path:
        Absolute path to the MCP server project on the host.  Used for
        read-only bind mounts.
    """
    config = ContainerConfig(image=image)

    # -- Network --
    _apply_network_policy(config, policy)

    # -- Filesystem --
    _apply_filesystem_policy(config, policy, project_path)

    # -- Resources --
    _apply_resource_limits(config, policy)

    # -- Labels --
    config.labels = {
        "hdx.sandbox": "true",
        "hdx.policy": policy.name,
    }

    return config


# ---------------------------------------------------------------------------
# Network mapping
# ---------------------------------------------------------------------------


def _apply_network_policy(config: ContainerConfig, policy: SandboxPolicy) -> None:
    """Set network_mode and domain filtering based on the network policy."""
    if not policy.network.enabled:
        # No network at all — strict isolation
        config.network_mode = "none"
        return

    if policy.network.allowed_domains:
        # Restricted network — use bridge with iptables-based domain filtering.
        # NET_ADMIN capability is required to set up firewall rules inside the
        # container after it starts.
        config.network_mode = "bridge"
        config.allowed_domains = list(policy.network.allowed_domains)
        config.cap_add = ["NET_ADMIN"]
    else:
        # Unrestricted network — permissive mode
        config.network_mode = "bridge"


# ---------------------------------------------------------------------------
# Filesystem mapping
# ---------------------------------------------------------------------------


def _apply_filesystem_policy(
    config: ContainerConfig,
    policy: SandboxPolicy,
    project_path: str,
) -> None:
    """Set volume mounts and tmpfs entries from the filesystem policy."""
    # Read-only bind mounts
    for mount_path in policy.filesystem.read_only_mounts:
        # Map host project_path to the container mount point
        host_path = project_path if project_path else mount_path
        config.volumes[host_path] = {"bind": mount_path, "mode": "ro"}

    # Writable directories as tmpfs (ephemeral — nothing persists to host)
    for writable_dir in policy.filesystem.writable_dirs:
        config.tmpfs[writable_dir] = "size=64m"


# ---------------------------------------------------------------------------
# Resource limits mapping
# ---------------------------------------------------------------------------


def _apply_resource_limits(config: ContainerConfig, policy: SandboxPolicy) -> None:
    """Set CPU, memory, and timeout constraints."""
    config.mem_limit = policy.resources.memory_limit
    config.timeout_seconds = policy.resources.timeout_seconds

    # Convert fractional CPU to cpu_period / cpu_quota.
    # cpu_period is the scheduling period in microseconds (default 100,000).
    # cpu_quota is the max time the container can use within that period.
    # E.g. cpu_limit="0.5" → quota=50,000 (half a core).
    cpu_period = 100_000
    try:
        cpu_fraction = float(policy.resources.cpu_limit)
    except (ValueError, TypeError):
        cpu_fraction = 1.0

    cpu_quota = int(cpu_period * cpu_fraction)
    # Clamp to reasonable bounds
    cpu_quota = max(1000, cpu_quota)

    config.cpu_period = cpu_period
    config.cpu_quota = cpu_quota


# ---------------------------------------------------------------------------
# Dockerfile generation
# ---------------------------------------------------------------------------

_PYTHON_DOCKERFILE_TEMPLATE = """\
FROM python:{python_version}-slim

WORKDIR /app

# Create non-root user for running the MCP server
RUN useradd -r -s /usr/sbin/nologin hdx
{system_packages}
# Install dependencies first for better layer caching
COPY requirements.txt* pyproject.toml* ./
RUN pip install --no-cache-dir -e . 2>/dev/null || \\
    ([ -f requirements.txt ] && pip install --no-cache-dir -r requirements.txt) || \\
    true

# Copy the rest of the project
COPY . .

# Install the project itself
RUN pip install --no-cache-dir -e . 2>/dev/null || true
{security_footer}"""

_IPTABLES_INSTALL = (
    "\n# Install iptables for domain allowlist enforcement\n"
    "RUN apt-get update && apt-get install -y --no-install-recommends iptables "
    "&& rm -rf /var/lib/apt/lists/*\n"
)


def _detect_python_version(project_path: str) -> str:
    """Detect the minimum Python version from the project's pyproject.toml.

    Returns a version string like ``"3.11"`` or ``"3.13"``.
    Falls back to ``"3.13"`` if detection fails.
    """
    pyproject = Path(project_path) / "pyproject.toml"
    if not pyproject.is_file():
        return "3.13"

    try:
        with open(pyproject, "rb") as f:
            data = tomllib.load(f)
    except Exception:
        return "3.13"

    requires_python = data.get("project", {}).get("requires-python", "")
    if not requires_python:
        return "3.13"

    # Extract the first version number (e.g. ">=3.11" → "3.11")
    match = re.search(r"(\d+\.\d+)", requires_python)
    if match:
        return match.group(1)

    return "3.13"


def generate_firewall_script(allowed_domains: list[str]) -> str:
    """Generate an iptables/ip6tables firewall script that restricts outbound traffic.

    The script:
    1. Allows loopback traffic (required for Docker's internal DNS at 127.0.0.11)
    2. Allows established/related connections
    3. Resolves each allowed domain and adds ACCEPT rules for those IPs
    4. Writes resolved IPs to ``/etc/hosts`` as a fallback for name resolution
    5. Sets default OUTPUT policy to DROP on both IPv4 and IPv6

    Both ``iptables`` (IPv4) and ``ip6tables`` (IPv6) are configured.
    Without IPv6 rules, applications bypass the firewall by connecting
    over IPv6.

    DNS resolution still works (Docker routes it through loopback) but
    actual connections to non-allowed IPs are blocked.  Applications can
    resolve ``google.com`` — they just can't connect to it.

    Wildcard domains (e.g. ``*.github.com``) are resolved using the base
    domain, which covers the common case but not all subdomains.
    """
    lines = [
        "#!/bin/sh",
        "",
        "# Resolve allowed domains first (while DNS is still available)",
    ]

    # Resolve all domains and collect IPs before touching any firewall rules.
    # This avoids partial rule application if resolution fails.
    for i, domain in enumerate(allowed_domains):
        resolve_domain = domain.lstrip("*.")
        var = f"RESOLVED_{i}"
        lines.append(f'{var}=$(getent ahosts "{resolve_domain}" '
                     "| awk '{print $1}' | sort -u)")
        lines.append(f'if [ -z "${var}" ]; then')
        lines.append(f'  echo "Warning: could not resolve {resolve_domain}" >&2')
        lines.append("fi")

    lines.extend([
        "",
        "# --- IPv4 rules (iptables) ---",
        "",
        "iptables -F OUTPUT",
        "iptables -A OUTPUT -o lo -j ACCEPT",
        "iptables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT",
    ])

    # Add IPv4 ACCEPT rules per domain
    for i, domain in enumerate(allowed_domains):
        resolve_domain = domain.lstrip("*.")
        var = f"RESOLVED_{i}"
        lines.append(f'for ip in ${var}; do')
        lines.append('  case "$ip" in')
        lines.append("    *:*) ;;  # skip IPv6")
        lines.append(f'    *) iptables -A OUTPUT -d "$ip" -j ACCEPT ;;')
        lines.append("  esac")
        lines.append("done")

    lines.extend([
        "",
        "# Set IPv4 default policy to DROP",
        "iptables -P OUTPUT DROP",
        "",
        "# --- IPv6 rules (ip6tables, best-effort) ---",
        "# ip6tables may not be available in all containers.",
        "",
        "if command -v ip6tables >/dev/null 2>&1; then",
        "  ip6tables -F OUTPUT",
        "  ip6tables -A OUTPUT -o lo -j ACCEPT",
        "  ip6tables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT",
    ])

    for i, domain in enumerate(allowed_domains):
        resolve_domain = domain.lstrip("*.")
        var = f"RESOLVED_{i}"
        lines.append(f'  for ip in ${var}; do')
        lines.append('    case "$ip" in')
        lines.append(f'      *:*) ip6tables -A OUTPUT -d "$ip" -j ACCEPT ;;')
        lines.append("      *) ;;  # skip IPv4")
        lines.append("    esac")
        lines.append("  done")

    lines.extend([
        "  ip6tables -P OUTPUT DROP",
        "fi",
        "",
        "# Write /etc/hosts entries for all resolved IPs",
    ])

    for i, domain in enumerate(allowed_domains):
        resolve_domain = domain.lstrip("*.")
        var = f"RESOLVED_{i}"
        lines.append(f'for ip in ${var}; do')
        lines.append(f'  echo "$ip {resolve_domain}" >> /etc/hosts')
        lines.append("done")
        lines.append("")

    return "\n".join(lines) + "\n"


def generate_dockerfile(
    project_path: str,
    *,
    network_filtering: bool = False,
    allowed_domains: list[str] | None = None,
) -> str | None:
    """Generate a minimal Dockerfile for a Python MCP server project.

    Returns the Dockerfile content as a string, or ``None`` if the project
    already has a Dockerfile.

    The generated Dockerfile:
    - Detects the Python version from ``requires-python`` in pyproject.toml
    - Copies and installs dependencies
    - Copies the full project source
    - When *allowed_domains* is provided, bakes the iptables firewall script
      and an entrypoint wrapper into the image so that domain filtering is
      enforced on every container start (survives stop/start cycles).
    """
    path = Path(project_path)

    # Don't overwrite an existing Dockerfile
    if (path / "Dockerfile").exists():
        return None

    # Infer network_filtering from allowed_domains if not explicitly set
    if allowed_domains:
        network_filtering = True

    python_version = _detect_python_version(project_path)
    system_packages = _IPTABLES_INSTALL if network_filtering else ""
    security_footer = _generate_security_footer(allowed_domains)

    return _PYTHON_DOCKERFILE_TEMPLATE.format(
        python_version=python_version,
        system_packages=system_packages,
        security_footer=security_footer,
    )


def _generate_security_footer(allowed_domains: list[str] | None = None) -> str:
    """Generate Dockerfile lines for privilege dropping and optional firewall.

    Always uses Docker's ``USER`` directive to run as the non-root ``hdx``
    user.  When domain filtering is active, the firewall script is baked into
    the image at ``/etc/hdx/firewall.sh`` — the engine executes it as root
    via ``docker exec --user root`` after the container starts, before running
    any tool calls.  The ``USER hdx`` directive ensures the main container
    process and all subsequent ``docker exec`` calls default to non-root.
    """
    if not allowed_domains:
        return "\n# Run as non-root user\nUSER hdx\n"

    firewall_script = generate_firewall_script(allowed_domains)
    firewall_b64 = base64.b64encode(firewall_script.encode()).decode()

    return (
        "\n# Embed firewall script for domain allowlist enforcement\n"
        f"RUN mkdir -p /etc/hdx \\\n"
        f"    && echo '{firewall_b64}' | base64 -d > /etc/hdx/firewall.sh \\\n"
        f"    && chmod +x /etc/hdx/firewall.sh\n"
        "\n"
        "# Run as non-root user\n"
        "USER hdx\n"
    )
