from __future__ import annotations

import argparse
import hashlib
import ipaddress
import json
import os
import re
import secrets
import subprocess
import uuid
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Final

ROOT: Final = Path(__file__).resolve().parent
ENV_PATH: Final = ROOT / ".env"
RUNTIME_ROOT: Final = ROOT / "runtime"
RUNTIME: Final = RUNTIME_ROOT / "mosquitto"
AGENT_REGISTRY_PATH: Final = RUNTIME_ROOT / "agents.json"
IMAGE: Final = "eclipse-mosquitto:2.1.2-alpine"
PLACEHOLDER: Final = "replace-with-"
AGENT_REGISTRY_SCHEMA: Final = "missioncore.telemetry-agent-credential-registry/v1"
AGENT_PAYLOAD_SCHEMA: Final = "missioncore.telemetry-agent-provisioning-payload/v1"
AGENT_BUNDLE_SCHEMA: Final = "missioncore.telemetry-agent-bundle/v1"
SAFE_IDENTIFIER: Final = re.compile(
    r"^[a-z0-9](?:[a-z0-9-]{0,62}[a-z0-9])?$"
)
SAFE_INTERVAL: Final = re.compile(r"^[1-9][0-9]{0,2}s$")
WINDOWS_BUNDLE_FILES: Final = (
    "Get-NdcMissionCorePipelineTelemetry.ps1",
    "Install-NdcMissionCoreTelegraf.ps1",
    "Update-NdcMissionCoreTelegraf.ps1",
    "mission-core-windows.conf.tmpl",
)


@dataclass(frozen=True, slots=True)
class AgentCredential:
    contour_id: str
    agent_id: str
    password: str

    def __post_init__(self) -> None:
        _safe_identifier(self.contour_id, "contour_id")
        _safe_identifier(self.agent_id, "agent_id")
        if (
            not 32 <= len(self.password) <= 512
            or self.password.startswith(PLACEHOLDER)
            or any(ord(character) < 33 for character in self.password)
        ):
            raise RuntimeError("agent password must contain a non-placeholder value")


def _initialize_environment(bind_address: str, *, overwrite: bool = False) -> None:
    if ENV_PATH.exists() and not overwrite:
        raise RuntimeError(".env already exists; refusing to overwrite local credentials")
    normalized_bind_address = str(ipaddress.ip_address(bind_address))
    values = {
        "MISSIONCORE_MQTT_BIND_ADDRESS": normalized_bind_address,
        "MISSIONCORE_MQTT_PORT": "1883",
        "MISSIONCORE_TELEMETRY_QUERY_PORT": "18030",
        "MISSIONCORE_DB_NAME": "missioncore_telemetry",
        "MISSIONCORE_DB_USER": "missioncore_ingest",
        "MISSIONCORE_DB_PASSWORD": secrets.token_urlsafe(36),
        "MISSIONCORE_DB_INGEST_PASSWORD": secrets.token_urlsafe(36),
        "MISSIONCORE_MQTT_INGEST_USER": "missioncore-ingest",
        "MISSIONCORE_MQTT_INGEST_PASSWORD": secrets.token_urlsafe(36),
    }
    _write_private(
        ENV_PATH,
        "".join(f"{name}={value}\n" for name, value in values.items()).encode(),
        replace=overwrite,
    )


def _environment() -> dict[str, str]:
    if not ENV_PATH.is_file():
        raise RuntimeError("run prepare.py --initialize and set unique secrets first")
    values: dict[str, str] = {}
    for raw_line in ENV_PATH.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        name, value = line.split("=", 1)
        values[name.strip()] = value.strip()
    return values


def _migrate_environment() -> None:
    values = _environment()
    additions: dict[str, str] = {}
    if "MISSIONCORE_DB_INGEST_PASSWORD" not in values:
        additions["MISSIONCORE_DB_INGEST_PASSWORD"] = secrets.token_urlsafe(36)
    if additions:
        with ENV_PATH.open("a", encoding="utf-8", newline="\n") as stream:
            for name, value in additions.items():
                stream.write(f"{name}={value}\n")
            stream.flush()
            os.fsync(stream.fileno())
        os.chmod(ENV_PATH, 0o600)
        values.update(additions)
    _migrate_legacy_worker(values)


def _migrate_legacy_worker(values: dict[str, str]) -> None:
    legacy_names = (
        "MISSIONCORE_MQTT_WORKER_006_USER",
        "MISSIONCORE_MQTT_WORKER_006_CONTOUR",
        "MISSIONCORE_MQTT_WORKER_006_PASSWORD",
    )
    if AGENT_REGISTRY_PATH.exists() or not all(values.get(name) for name in legacy_names):
        return
    credential = AgentCredential(
        contour_id=_identifier(values, legacy_names[1]),
        agent_id=_identifier(values, legacy_names[0]),
        password=_required(values, legacy_names[2]),
    )
    _write_agent_registry((credential,))


def _read_agent_registry(*, missing_ok: bool = False) -> tuple[AgentCredential, ...]:
    if not AGENT_REGISTRY_PATH.exists():
        if missing_ok:
            return ()
        raise RuntimeError("no telemetry agents enrolled; use --enroll-agent first")
    if AGENT_REGISTRY_PATH.is_symlink() or not AGENT_REGISTRY_PATH.is_file():
        raise RuntimeError("telemetry agent registry must be a regular file")
    if AGENT_REGISTRY_PATH.stat().st_size > 1024 * 1024:
        raise RuntimeError("telemetry agent registry is too large")
    try:
        payload: object = json.loads(AGENT_REGISTRY_PATH.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise RuntimeError("telemetry agent registry is unreadable") from exc
    if not isinstance(payload, dict) or set(payload) != {"schema_version", "agents"}:
        raise RuntimeError("telemetry agent registry shape is invalid")
    if payload["schema_version"] != AGENT_REGISTRY_SCHEMA:
        raise RuntimeError("telemetry agent registry schema is invalid")
    rows = payload["agents"]
    if not isinstance(rows, list):
        raise RuntimeError("telemetry agent registry agents must be an array")
    credentials: list[AgentCredential] = []
    for row in rows:
        if not isinstance(row, dict) or set(row) != {"contour_id", "agent_id", "password"}:
            raise RuntimeError("telemetry agent registry row is invalid")
        if not all(isinstance(value, str) for value in row.values()):
            raise RuntimeError("telemetry agent registry values must be strings")
        credentials.append(
            AgentCredential(
                contour_id=row["contour_id"],
                agent_id=row["agent_id"],
                password=row["password"],
            )
        )
    _reject_agent_duplicates(tuple(credentials))
    os.chmod(AGENT_REGISTRY_PATH, 0o600)
    return tuple(credentials)


def _write_agent_registry(credentials: tuple[AgentCredential, ...]) -> None:
    _reject_agent_duplicates(credentials)
    document = {
        "schema_version": AGENT_REGISTRY_SCHEMA,
        "agents": [
            {
                "contour_id": credential.contour_id,
                "agent_id": credential.agent_id,
                "password": credential.password,
            }
            for credential in sorted(credentials, key=lambda item: item.agent_id)
        ],
    }
    _write_private(
        AGENT_REGISTRY_PATH,
        _canonical_json(document) + b"\n",
        replace=True,
    )


def _enroll_agent(contour_id: str, agent_id: str) -> AgentCredential:
    credential = AgentCredential(
        contour_id=contour_id,
        agent_id=agent_id,
        password=secrets.token_urlsafe(36),
    )
    existing = _read_agent_registry(missing_ok=True)
    if any(row.agent_id == agent_id for row in existing):
        raise RuntimeError(f"telemetry agent is already enrolled: {agent_id}")
    _write_agent_registry((*existing, credential))
    return credential


def _export_agent_payload(
    *,
    agent_id: str,
    node_id: str,
    output: Path,
    mqtt_host: str | None = None,
    telemetry_interval: str = "2s",
) -> None:
    if (
        not node_id
        or node_id != node_id.strip()
        or len(node_id) > 256
        or any(ord(character) < 32 for character in node_id)
    ):
        raise RuntimeError("node_id is invalid")
    if SAFE_INTERVAL.fullmatch(telemetry_interval) is None:
        raise RuntimeError("telemetry interval must be 1-999 seconds")
    matches = [row for row in _read_agent_registry() if row.agent_id == agent_id]
    if len(matches) != 1:
        raise RuntimeError(f"telemetry agent is not enrolled: {agent_id}")
    credential = matches[0]
    values = _environment()
    resolved_host = mqtt_host or _required(values, "MISSIONCORE_MQTT_BIND_ADDRESS")
    _private_mqtt_host(resolved_host)
    port = _required(values, "MISSIONCORE_MQTT_PORT")
    if not port.isdigit() or not 1 <= int(port) <= 65535:
        raise RuntimeError("MISSIONCORE_MQTT_PORT is invalid")
    payload = {
        "schema_version": AGENT_PAYLOAD_SCHEMA,
        "MISSIONCORE_CONTOUR_ID": credential.contour_id,
        "MISSIONCORE_AGENT_ID": credential.agent_id,
        "MISSIONCORE_NODE_ID": node_id,
        "MISSIONCORE_MQTT_HOST": resolved_host,
        "MISSIONCORE_MQTT_PORT": port,
        "MISSIONCORE_MQTT_USERNAME": credential.agent_id,
        "MISSIONCORE_MQTT_PASSWORD": credential.password,
        "MISSIONCORE_TELEMETRY_INTERVAL": telemetry_interval,
    }
    _write_private(output, _canonical_json(payload) + b"\n", replace=False)


def _build_agent_bundle(*, platform_name: str, output: Path) -> str:
    if platform_name != "windows":
        raise RuntimeError("only the reviewed windows telemetry agent bundle is available")
    candidate = output.expanduser().absolute()
    if candidate.exists() or candidate.is_symlink():
        raise RuntimeError("agent bundle output already exists")
    sources = [(name, ROOT / "telegraf" / name) for name in WINDOWS_BUNDLE_FILES]
    inventory = []
    for name, path in sources:
        if path.is_symlink() or not path.is_file():
            raise RuntimeError(f"agent bundle source is unavailable: {name}")
        content = path.read_bytes()
        inventory.append(
            {
                "path": name,
                "byte_length": len(content),
                "sha256": hashlib.sha256(content).hexdigest(),
            }
        )
    identity = {
        "schema_version": AGENT_BUNDLE_SCHEMA,
        "platform": platform_name,
        "credential_embedded": False,
        "files": inventory,
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    manifest = {
        **identity,
        "bundle_id": f"missioncore-telemetry-agent-{platform_name}-{identity_sha256}",
        "identity_sha256": identity_sha256,
    }
    candidate.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    staging = candidate.parent / f".{candidate.name}.{uuid.uuid4().hex}.tmp"
    with zipfile.ZipFile(
        staging,
        "w",
        compression=zipfile.ZIP_DEFLATED,
        compresslevel=9,
    ) as archive:
        _zip_entry(archive, "manifest.json", _canonical_json(manifest) + b"\n")
        for name, path in sources:
            _zip_entry(archive, name, path.read_bytes())
    os.chmod(staging, 0o600)
    os.replace(staging, candidate)
    return str(manifest["bundle_id"])


def _required(values: dict[str, str], name: str) -> str:
    value = values.get(name, "")
    if not value or value.startswith(PLACEHOLDER):
        raise RuntimeError(f"{name} must contain a non-placeholder value")
    return value


def _identifier(values: dict[str, str], name: str) -> str:
    return _safe_identifier(_required(values, name), name)


def _safe_identifier(value: str, label: str) -> str:
    if SAFE_IDENTIFIER.fullmatch(value) is None:
        raise RuntimeError(f"{label} must contain a DNS-safe lowercase identifier")
    return value


def _private_mqtt_host(value: str) -> None:
    try:
        address = ipaddress.ip_address(value)
    except ValueError:
        if not value.endswith(".local") or len(value) > 253 or any(
            SAFE_IDENTIFIER.fullmatch(label) is None for label in value[:-6].split(".")
        ):
            raise RuntimeError(
                "MQTT host must be a private IP or stable .local hostname"
            ) from None
    else:
        if not address.is_private or address.is_unspecified or address.is_multicast:
            raise RuntimeError("MQTT host must be a private IP or stable .local hostname")


def _reject_agent_duplicates(credentials: tuple[AgentCredential, ...]) -> None:
    agent_ids = [credential.agent_id for credential in credentials]
    if len(agent_ids) != len(set(agent_ids)):
        raise RuntimeError("telemetry agent_id must be globally unique")


def _password_entry(path: Path, username: str, password: str, *, create: bool) -> None:
    command = [
        "docker",
        "run",
        "--rm",
        "-i",
        "-v",
        f"{RUNTIME}:/out",
        IMAGE,
        "mosquitto_passwd",
    ]
    if create:
        command.append("-c")
    command.extend([f"/out/{path.name}", username])
    subprocess.run(
        command,
        check=True,
        input=f"{password}\n{password}\n",
        text=True,
    )


def _prepare_password_entries(
    path: Path,
    ingest_user: str,
    ingest_password: str,
    agents: tuple[AgentCredential, ...],
) -> None:
    staging = path.parent / f".{path.name}.{uuid.uuid4().hex}.tmp"
    try:
        _password_entry(staging, ingest_user, ingest_password, create=True)
        for credential in sorted(agents, key=lambda item: item.agent_id):
            _password_entry(
                staging,
                credential.agent_id,
                credential.password,
                create=False,
            )
        os.chmod(staging, 0o600)
        os.replace(staging, path)
    finally:
        if staging.exists():
            staging.unlink()


def _acl_document(ingest_user: str, agents: tuple[AgentCredential, ...]) -> str:
    lines = [
        f"user {ingest_user}",
        "topic read mission-core/v1/contours/+/agents/+/+",
        "topic read $SYS/broker/uptime",
        "",
        "# Every credential is scoped to one contour and one stable agent.",
    ]
    for credential in sorted(agents, key=lambda item: item.agent_id):
        lines.extend(
            [
                f"user {credential.agent_id}",
                (
                    "topic write mission-core/v1/contours/"
                    f"{credential.contour_id}/agents/{credential.agent_id}/+"
                ),
                "",
            ]
        )
    return "\n".join(lines)


def _prepare_plane() -> None:
    values = _environment()
    _migrate_legacy_worker(values)
    agents = _read_agent_registry()
    ingest_user = _required(values, "MISSIONCORE_MQTT_INGEST_USER")
    ingest_password = _required(values, "MISSIONCORE_MQTT_INGEST_PASSWORD")
    _required(values, "MISSIONCORE_DB_PASSWORD")
    RUNTIME.mkdir(mode=0o700, parents=True, exist_ok=True)
    password_path = RUNTIME / "passwords"
    _prepare_password_entries(password_path, ingest_user, ingest_password, agents)
    acl_path = RUNTIME / "acl"
    _write_private(
        acl_path,
        _acl_document(ingest_user, agents).encode("utf-8"),
        replace=True,
    )


def _write_private(path: Path, content: bytes, *, replace: bool) -> None:
    candidate = path.expanduser().absolute()
    if candidate.is_symlink() or (candidate.exists() and not replace):
        raise RuntimeError(f"refusing to replace private file: {candidate.name}")
    candidate.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    staging = candidate.parent / f".{candidate.name}.{uuid.uuid4().hex}.tmp"
    descriptor = os.open(staging, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
    try:
        os.write(descriptor, content)
        os.fsync(descriptor)
    finally:
        os.close(descriptor)
    try:
        os.replace(staging, candidate)
        os.chmod(candidate, 0o600)
    finally:
        if staging.exists():
            staging.unlink()


def _zip_entry(archive: zipfile.ZipFile, name: str, content: bytes) -> None:
    info = zipfile.ZipInfo(name, date_time=(2026, 1, 1, 0, 0, 0))
    info.compress_type = zipfile.ZIP_DEFLATED
    info.external_attr = 0o600 << 16
    info.create_system = 3
    archive.writestr(info, content, compress_type=zipfile.ZIP_DEFLATED, compresslevel=9)


def _canonical_json(value: object) -> bytes:
    return json.dumps(
        value,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    ).encode("utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--initialize", action="store_true")
    parser.add_argument("--migrate", action="store_true")
    parser.add_argument("--mqtt-bind-address", default="127.0.0.1")
    parser.add_argument("--enroll-agent", metavar="AGENT_ID")
    parser.add_argument("--contour-id")
    parser.add_argument("--export-agent-payload", metavar="AGENT_ID")
    parser.add_argument("--node-id")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--mqtt-host")
    parser.add_argument("--telemetry-interval", default="2s")
    parser.add_argument("--build-agent-bundle", choices=("windows",))
    arguments = parser.parse_args()

    action_count = sum(
        bool(value)
        for value in (
            arguments.initialize,
            arguments.migrate,
            arguments.enroll_agent,
            arguments.export_agent_payload,
            arguments.build_agent_bundle,
        )
    )
    if action_count > 1:
        raise RuntimeError("select exactly one provisioning action")
    if arguments.initialize:
        _initialize_environment(arguments.mqtt_bind_address)
        print("Private telemetry-plane environment initialized.", flush=True)
        return
    if arguments.migrate:
        _migrate_environment()
        print("Private telemetry-plane state migrated.", flush=True)
        return
    if arguments.enroll_agent:
        if not arguments.contour_id:
            raise RuntimeError("--contour-id is required with --enroll-agent")
        credential = _enroll_agent(arguments.contour_id, arguments.enroll_agent)
        print(
            f"Telemetry agent enrolled: {credential.agent_id} ({credential.contour_id}).",
            flush=True,
        )
        return
    if arguments.export_agent_payload:
        if not arguments.node_id or arguments.output is None:
            raise RuntimeError("--node-id and --output are required for payload export")
        _export_agent_payload(
            agent_id=arguments.export_agent_payload,
            node_id=arguments.node_id,
            output=arguments.output,
            mqtt_host=arguments.mqtt_host,
            telemetry_interval=arguments.telemetry_interval,
        )
        print(f"Private agent payload written: {arguments.output.name}.", flush=True)
        return
    if arguments.build_agent_bundle:
        if arguments.output is None:
            raise RuntimeError("--output is required for agent bundle build")
        bundle_id = _build_agent_bundle(
            platform_name=arguments.build_agent_bundle,
            output=arguments.output,
        )
        print(f"Secret-free agent bundle built: {bundle_id}.", flush=True)
        return
    _prepare_plane()
    print("Mosquitto password and ACL files prepared for all enrolled agents.", flush=True)


if __name__ == "__main__":
    main()