#!/usr/bin/env python3
import os
import socket
import subprocess
import uuid

from flask import Flask, request, render_template_string, flash

app = Flask(__name__)
app.secret_key = os.environ.get("MICROCONTAINERS_SECRET_KEY", "dev-only-not-secret")

# --- Config -----------------------------------------------------------------

IMAGES = {
    "debian": {
        "docker_image": "micro-debian-dev:latest",
        "label": "Debian Dev Box (micro-debian-dev)"
    },
    "fedora": {
        "docker_image": "micro-fedora43-dev:latest",
        "label": "Fedora 43 Dev Box (micro-fedora43-dev)"
    },
}

DEFAULT_IMAGE_KEY = "debian"

DATA_ROOT = "/srv/microcontainers"         # per-tenant metadata + disk.img
MOUNT_ROOT = "/mnt/microcontainers"        # where disk.img files are mounted
HOSTNAME = "arthur.lan"                    # or "mentalnet.xyz" if you prefer
PORT_RANGE_START = 20000
PORT_RANGE_END = 21000

os.makedirs(DATA_ROOT, exist_ok=True)
os.makedirs(MOUNT_ROOT, exist_ok=True)

# --- Helpers ----------------------------------------------------------------


def find_free_port(start=PORT_RANGE_START, end=PORT_RANGE_END):
    """Find a free TCP port on the host in the given range."""
    for port in range(start, end):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(0.2)
            try:
                s.bind(("", port))
                return port
            except OSError:
                continue
    raise RuntimeError("No free ports available in range")


def validate_ssh_key(key: str) -> bool:
    """Very basic SSH public key validation."""
    key = key.strip()
    if not key:
        return False
    allowed_prefixes = (
        "ssh-ed25519",
        "ssh-rsa",
        "ecdsa-sha2-",
        "sk-ssh-",
    )
    return any(key.startswith(p) for p in allowed_prefixes)


def is_mounted(mount_point: str) -> bool:
    """Check if a given mount point is already mounted."""
    try:
        with open("/proc/mounts", "r") as f:
            for line in f:
                parts = line.split()
                if len(parts) >= 2 and parts[1] == mount_point:
                    return True
    except FileNotFoundError:
        pass
    return False


def ensure_tenant_disk(container_name: str, ssh_key: str, size_gb: int = 3) -> str:
    """
    Ensure a per-tenant disk.img exists, is formatted, mounted, and has
    .ssh/authorized_keys populated.

    Returns the host mount point path, which will be bind-mounted to /home/micro.
    """
    tenant_root = os.path.join(DATA_ROOT, container_name)
    os.makedirs(tenant_root, exist_ok=True)

    disk_img = os.path.join(tenant_root, "disk.img")
    mount_point = os.path.join(MOUNT_ROOT, container_name)
    os.makedirs(mount_point, exist_ok=True)

    # 1) Create sparse disk.img if it doesn't exist
    if not os.path.exists(disk_img):
        subprocess.run(["fallocate", "-l", f"{size_gb}G", disk_img], check=True)
        subprocess.run(["mkfs.ext4", "-F", disk_img], check=True)

    # 2) Mount it if not already mounted
    if not is_mounted(mount_point):
        subprocess.run(["mount", "-o", "loop", disk_img, mount_point], check=True)

    # 3) Prepare /home/micro layout inside the filesystem
    #    (we're mounting this whole thing as /home/micro in the container)
    ssh_dir = os.path.join(mount_point, ".ssh")
    os.makedirs(ssh_dir, exist_ok=True)

    auth_keys_path = os.path.join(ssh_dir, "authorized_keys")
    with open(auth_keys_path, "w") as f:
        f.write(ssh_key.strip() + "\n")

    # Proper perms for SSH
    os.chmod(ssh_dir, 0o700)
    os.chmod(auth_keys_path, 0o600)

    # Make everything in /home/micro look like micro:micro (uid/gid 1000)
    subprocess.run(["chown", "-R", "1000:1000", mount_point], check=True)

    return mount_point


def run_docker_container(container_name: str, host_port: int, docker_image: str, home_mount: str):
    """
    Start the tenant container:
      - Binds host_port -> container:22
      - Mounts per-tenant 3GB disk as /home/micro
    """
    cmd = [
        "docker", "run", "-d",
        "--name", container_name,
        "--memory=512m",
        "--memory-swap=512m",
        "-p", f"{host_port}:22",
        "-v", f"{home_mount}:/home/micro",
        docker_image,
    ]
    subprocess.run(cmd, check=True)


# --- Template ---------------------------------------------------------------

INDEX_TEMPLATE = """
<!doctype html>
<html>
<head>
    <meta charset="utf-8">
    <title>MicroContainer Provisioning</title>
    <style>
        body {
            background: #111;
            color: #eee;
            font-family: system-ui, -apple-system, BlinkMacSystemFont, sans-serif;
            max-width: 720px;
            margin: 40px auto;
            padding: 0 20px;
        }
        h1 { color: #66e0ff; }
        .box {
            background: #1a1a1a;
            border-radius: 8px;
            padding: 20px;
            margin-bottom: 20px;
            border: 1px solid #333;
        }
        label { display: block; margin-bottom: 8px; font-weight: 600; }
        textarea, input[type=text], select {
            width: 100%;
            box-sizing: border-box;
            background: #101010;
            border: 1px solid #444;
            border-radius: 4px;
            color: #eee;
            padding: 8px;
            font-family: monospace;
            font-size: 0.9rem;
        }
        button {
            background: #66e0ff;
            color: #000;
            padding: 8px 16px;
            border-radius: 4px;
            border: none;
            font-weight: 600;
            cursor: pointer;
            margin-top: 10px;
        }
        button:hover { background: #4bc1dd; }
        .message { color: #ff8080; margin-bottom: 10px; }
        .success { color: #9bff9b; }
        code { background: #000; padding: 2px 4px; border-radius: 3px; }
    </style>
</head>
<body>
    <h1>MicroContainer Provisioning (Prototype)</h1>

    <div class="box">
        <p>This prototype spins up a small Docker-based "micro VM" and wires in your SSH public key.</p>
        <ul>
            <li>Available base images:</li>
            <ul>
              {% for key, info in images.items() %}
                <li><code>{{ key }}</code>: {{ info.label }}</li>
              {% endfor %}
            </ul>
            <li>User: <code>micro</code></li>
            <li>Per-tenant /home size: <code>3 GB</code> (loopback ext4)</li>
            <li>SSH inside container: <code>port 22</code></li>
            <li>Host port: auto-assigned from <code>{{ port_start }}-{{ port_end }}</code></li>
        </ul>
    </div>

    {% with messages = get_flashed_messages() %}
      {% if messages %}
        <div class="message">
          {% for msg in messages %}
            {{ msg }}<br>
          {% endfor %}
        </div>
      {% endif %}
    {% endwith %}

    {% if result %}
    <div class="box">
        <h2 class="success">Container created!</h2>
        <p><strong>Image:</strong> <code>{{ result.image_key }}</code> — {{ images[result.image_key].label }}</p>
        <p><strong>Container name:</strong> <code>{{ result.name }}</code></p>
        <p><strong>Host port:</strong> <code>{{ result.port }}</code></p>
        <p><strong>SSH command:</strong></p>
        <pre><code>ssh micro@{{ hostname }} -p {{ result.port }}</code></pre>
        <p>You can change the password for <code>micro</code> after logging in with:</p>
        <pre><code>passwd</code></pre>
    </div>
    {% endif %}

    <div class="box">
        <form method="post">
            <label for="image">Base image</label>
            <select id="image" name="image">
              {% for key, info in images.items() %}
                <option value="{{ key }}" {% if key == selected_image %}selected{% endif %}>
                  {{ key }} - {{ info.label }}
                </option>
              {% endfor %}
            </select>

            <label for="ssh_key">SSH public key</label>
            <textarea id="ssh_key" name="ssh_key" rows="6" placeholder="ssh-ed25519 AAAAC3Nz... user@host" required>{{ ssh_key or "" }}</textarea>

            <label for="note">Optional note / label (will be part of container name)</label>
            <input type="text" id="note" name="note" placeholder="e.g. debian-dev-01">

            <button type="submit">Provision MicroContainer</button>
        </form>
    </div>
</body>
</html>
"""


# --- Routes -----------------------------------------------------------------


@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "GET":
        return render_template_string(
            INDEX_TEMPLATE,
            images=IMAGES,
            port_start=PORT_RANGE_START,
            port_end=PORT_RANGE_END,
            result=None,
            ssh_key="",
            hostname=HOSTNAME,
            selected_image=DEFAULT_IMAGE_KEY,
        )

    # POST: handle provisioning
    ssh_key = (request.form.get("ssh_key") or "").strip()
    note = (request.form.get("note") or "").strip()
    image_key = (request.form.get("image") or DEFAULT_IMAGE_KEY).strip()

    if image_key not in IMAGES:
        image_key = DEFAULT_IMAGE_KEY

    docker_image = IMAGES[image_key]["docker_image"]

    if not validate_ssh_key(ssh_key):
        flash("Invalid SSH public key format. Please paste a standard OpenSSH public key line.")
        return render_template_string(
            INDEX_TEMPLATE,
            images=IMAGES,
            port_start=PORT_RANGE_START,
            port_end=PORT_RANGE_END,
            result=None,
            ssh_key=ssh_key,
            hostname=HOSTNAME,
            selected_image=image_key,
        )

    # Generate a container name
    suffix = uuid.uuid4().hex[:6]
    base_name = f"mc-{image_key}"
    if note:
        safe_note = "".join(c for c in note.lower().replace(" ", "-") if c.isalnum() or c in "-_")
        base_name = f"mc-{image_key}-{safe_note}"
    container_name = f"{base_name}-{suffix}"

    # Ensure per-tenant disk (3 GB ext4) and authorized_keys
    try:
        home_mount = ensure_tenant_disk(container_name, ssh_key, size_gb=3)
    except subprocess.CalledProcessError as e:
        flash(f"Error preparing tenant disk: {e}")
        return render_template_string(
            INDEX_TEMPLATE,
            images=IMAGES,
            port_start=PORT_RANGE_START,
            port_end=PORT_RANGE_END,
            result=None,
            ssh_key=ssh_key,
            hostname=HOSTNAME,
            selected_image=image_key,
        )

    # Choose a free port and start container
    host_port = find_free_port()
    try:
        run_docker_container(container_name, host_port, docker_image, home_mount)
    except subprocess.CalledProcessError as e:
        flash(f"Error starting container: {e}")
        return render_template_string(
            INDEX_TEMPLATE,
            images=IMAGES,
            port_start=PORT_RANGE_START,
            port_end=PORT_RANGE_END,
            result=None,
            ssh_key=ssh_key,
            hostname=HOSTNAME,
            selected_image=image_key,
        )

    class Result:
        def __init__(self, name, port, image_key):
            self.name = name
            self.port = port
            self.image_key = image_key

    result = Result(container_name, host_port, image_key)

    return render_template_string(
        INDEX_TEMPLATE,
        images=IMAGES,
        port_start=PORT_RANGE_START,
        port_end=PORT_RANGE_END,
        result=result,
        ssh_key="",
        hostname=HOSTNAME,
        selected_image=image_key,
    )


if __name__ == "__main__":
    # For dev/testing only
    app.run(host="0.0.0.0", port=5000, debug=True)