#!/usr/bin/env python3
"""
ZKAC opaque file-share server (demo).

Headless TCP service that combines:

* The same management-channel wire protocol as ``zkac-node serve`` so existing
  ``zkac-node registry create/get/update`` commands work unchanged.
* A new file-share channel that, after BBS+ role authentication, exposes
  bucket primitives. The server only ever sees opaque ciphertext blobs and
  per-recipient encrypted role grants; file names, contents and per-role
  visibility masks are never visible server-side.

Run with::

    uv run python demo/file_share_server.py --host 127.0.0.1 --port 9879

The server transport public key is printed at start-up; pin it on each client
with ``zkac-node server pin <userid> 127.0.0.1:9879 --key <hex>``.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import socket
import sys
import threading
import traceback
import uuid
from pathlib import Path
from typing import Any

import zkac
from zkac.tcp import FramedSession, server_handshake_anon


# ── small helpers ────────────────────────────────────────────────────

def _b64(data: bytes) -> str:
    return base64.b64encode(data).decode("ascii")


def _unb64(s: str) -> bytes:
    return base64.b64decode(s)


def _chmod(path: Path, mode: int) -> None:
    try:
        os.chmod(path, mode)
    except OSError:
        pass


def _write_private_json(path: Path, payload: dict) -> None:
    path.write_text(json.dumps(payload, indent=2))
    _chmod(path, 0o600)


def _is_loopback(host: str) -> bool:
    return host.strip().lower() in {"127.0.0.1", "::1", "localhost"}


def _safe_id(value: str) -> str:
    """Strict allowlist on identifiers used as filesystem names."""
    if not isinstance(value, str) or not value:
        raise ValueError("missing identifier")
    if len(value) > 128 or any(c not in "0123456789abcdefABCDEF-" for c in value):
        raise ValueError(f"invalid identifier {value!r}")
    return value


# ── opaque on-disk store ─────────────────────────────────────────────

class _FileShareStore:
    """Persists registry snapshots and opaque bucket state under one data dir."""

    def __init__(self, data_dir: Path) -> None:
        self._dir = data_dir
        self._reg_dir = data_dir / "registries"
        self._buckets_dir = data_dir / "buckets"
        for d in (self._dir, self._reg_dir, self._buckets_dir):
            d.mkdir(parents=True, exist_ok=True)
            _chmod(d, 0o700)
        self._lock = threading.Lock()

    # transport key

    def load_or_create_keypair(self) -> zkac.Keypair:
        kf = self._dir / "server_key.json"
        if kf.exists():
            data = json.loads(kf.read_text())
            return zkac.Keypair.from_secret_key(_unb64(data["secret_b64"]))
        kp = zkac.Keypair()
        _write_private_json(kf, {
            "secret_b64": _b64(kp.secret_key_bytes()),
            "public_b64": _b64(kp.public_key().to_bytes()),
        })
        return kp

    # registries (mirrors zkac-node serve so the CLI works unchanged)

    def save_registry(self, rid_hex: str, state_bytes: bytes, cert_bytes: bytes) -> None:
        with self._lock:
            (self._reg_dir / f"{rid_hex}.state").write_bytes(state_bytes)
            (self._reg_dir / f"{rid_hex}.cert").write_bytes(cert_bytes)

    def load_all_registries(self, mgr: zkac.RegistryManager) -> int:
        n = 0
        for p in sorted(self._reg_dir.glob("*.state")):
            cert = self._reg_dir / f"{p.stem}.cert"
            if not cert.exists():
                continue
            try:
                mgr.restore(p.read_bytes(), cert.read_bytes())
                n += 1
            except Exception as exc:
                print(f"[fs-server] skip registry {p.stem}: {exc}")
        return n

    # buckets

    def _bucket_dir(self, bucket_id: str) -> Path:
        return self._buckets_dir / _safe_id(bucket_id)

    def _bucket_meta_path(self, bucket_id: str) -> Path:
        return self._bucket_dir(bucket_id) / "meta.json"

    def bucket_create(self, bucket_id: str, owner_registry_id_hex: str) -> None:
        bd = self._bucket_dir(bucket_id)
        with self._lock:
            if bd.exists():
                raise RuntimeError("bucket already exists")
            (bd / "blobs").mkdir(parents=True)
            (bd / "role_grants").mkdir()
            _chmod(bd, 0o700)
            _write_private_json(self._bucket_meta_path(bucket_id), {
                "bucket_id": bucket_id,
                "owner_registry_id": owner_registry_id_hex,
                "finalized": False,
                "role_acl": {},  # role_id_hex -> {"version": int, "allowed_blob_ids": [blob_id...]}
            })

    def bucket_meta(self, bucket_id: str) -> dict:
        return json.loads(self._bucket_meta_path(bucket_id).read_text())

    def _bucket_require_owner(self, bucket_id: str, registry_id_hex: str) -> dict:
        meta = self.bucket_meta(bucket_id)
        if meta.get("owner_registry_id") != registry_id_hex:
            raise RuntimeError("not bucket owner")
        return meta

    def bucket_set_finalized(self, bucket_id: str, registry_id_hex: str, finalized: bool) -> None:
        with self._lock:
            meta = self._bucket_require_owner(bucket_id, registry_id_hex)
            meta["finalized"] = bool(finalized)
            _write_private_json(self._bucket_meta_path(bucket_id), meta)

    def bucket_delete(self, bucket_id: str, registry_id_hex: str) -> None:
        with self._lock:
            self._bucket_require_owner(bucket_id, registry_id_hex)
            bd = self._bucket_dir(bucket_id)
            for sub in ("blobs", "role_grants"):
                d = bd / sub
                if d.is_dir():
                    for p in d.iterdir():
                        try:
                            p.unlink()
                        except OSError:
                            pass
                    try:
                        d.rmdir()
                    except OSError:
                        pass
            try:
                self._bucket_meta_path(bucket_id).unlink()
            except OSError:
                pass
            try:
                bd.rmdir()
            except OSError:
                pass

    def bucket_put_blob(
        self,
        bucket_id: str,
        registry_id_hex: str,
        blob_id: str,
        ciphertext: bytes,
    ) -> None:
        with self._lock:
            self._bucket_require_owner(bucket_id, registry_id_hex)
            (self._bucket_dir(bucket_id) / "blobs" / _safe_id(blob_id)).write_bytes(ciphertext)

    def bucket_put_role_grant(
        self,
        bucket_id: str,
        registry_id_hex: str,
        recipient_pk_hex: str,
        role_id_hex: str,
        acl_version: int,
        eph_pk_b64: str,
        ciphertext_b64: str,
    ) -> None:
        with self._lock:
            self._bucket_require_owner(bucket_id, registry_id_hex)
            _safe_id(recipient_pk_hex)
            payload = {
                "bucket_id": bucket_id,
                "role_id_hex": _safe_id(role_id_hex),
                "acl_version": int(acl_version),
                "eph_pk_b64": eph_pk_b64,
                "ciphertext_b64": ciphertext_b64,
            }
            target = self._bucket_dir(bucket_id) / "role_grants" / f"{recipient_pk_hex}.json"
            _write_private_json(target, payload)

    def bucket_set_role_acl(
        self,
        bucket_id: str,
        registry_id_hex: str,
        role_id_hex: str,
        allowed_blob_ids: list[str],
    ) -> None:
        with self._lock:
            meta = self._bucket_require_owner(bucket_id, registry_id_hex)
            role_acl = dict(meta.get("role_acl", {}))
            key = _safe_id(role_id_hex)
            prev = role_acl.get(key, {})
            prev_version = int(prev.get("version", 0)) if isinstance(prev, dict) else 0
            role_acl[key] = {
                "version": prev_version + 1,
                "allowed_blob_ids": [_safe_id(b) for b in allowed_blob_ids],
            }
            meta["role_acl"] = role_acl
            _write_private_json(self._bucket_meta_path(bucket_id), meta)

    def bucket_get_blob(self, bucket_id: str, blob_id: str) -> bytes:
        return (self._bucket_dir(bucket_id) / "blobs" / _safe_id(blob_id)).read_bytes()

    def bucket_blob_allowed_for_role(self, bucket_id: str, role_id_hex: str, blob_id: str) -> bool:
        meta = self.bucket_meta(bucket_id)
        role_acl = meta.get("role_acl", {})
        role_meta = role_acl.get(_safe_id(role_id_hex))
        if not isinstance(role_meta, dict):
            return False
        allowed = role_meta.get("allowed_blob_ids")
        if not isinstance(allowed, list):
            return False
        return _safe_id(blob_id) in {str(v) for v in allowed}

    def bucket_role_acl(self, bucket_id: str, role_id_hex: str) -> dict:
        meta = self.bucket_meta(bucket_id)
        role_acl = meta.get("role_acl", {})
        role_meta = role_acl.get(_safe_id(role_id_hex), {})
        if not isinstance(role_meta, dict):
            return {"version": 0, "allowed_blob_ids": []}
        version = int(role_meta.get("version", 0))
        allowed = role_meta.get("allowed_blob_ids", [])
        if not isinstance(allowed, list):
            allowed = []
        return {
            "version": version,
            "allowed_blob_ids": [_safe_id(str(v)) for v in allowed],
        }

    def bucket_get_role_grant(self, bucket_id: str, recipient_pk_hex: str) -> dict:
        path = self._bucket_dir(bucket_id) / "role_grants" / f"{_safe_id(recipient_pk_hex)}.json"
        return json.loads(path.read_text())

    def buckets_for_recipient(self, recipient_pk_hex: str) -> list[str]:
        """Bucket ids that have an encrypted role grant addressed to ``recipient_pk_hex``."""
        _safe_id(recipient_pk_hex)
        out: list[str] = []
        for bd in sorted(self._buckets_dir.iterdir()):
            grant = bd / "role_grants" / f"{recipient_pk_hex}.json"
            if grant.is_file():
                out.append(bd.name)
        return out

    def buckets_owned_by(self, registry_id_hex: str) -> list[str]:
        out: list[str] = []
        for bd in sorted(self._buckets_dir.iterdir()):
            meta = bd / "meta.json"
            if not meta.is_file():
                continue
            try:
                if json.loads(meta.read_text()).get("owner_registry_id") == registry_id_hex:
                    out.append(bd.name)
            except (OSError, json.JSONDecodeError):
                continue
        return out


# ── command dispatch (inside encrypted session) ──────────────────────

def _dispatch_mgmt(
    cmd: dict,
    mgr: zkac.RegistryManager,
    store: _FileShareStore,
    server_pk_b64: str,
    transcript_hash: bytes,
) -> dict:
    """Registry management commands, wire-compatible with ``zkac-node`` CLI."""
    try:
        action = cmd.get("cmd")
        rid_hex = cmd.get("auth_registry_id")
        admin_proof_b64 = cmd.get("admin_proof_b64")

        def _require_admin_for(target_rid_hex: str) -> None:
            if rid_hex != target_rid_hex:
                raise RuntimeError("auth_registry_id must match command registry_id")
            if not isinstance(admin_proof_b64, str) or not admin_proof_b64:
                raise RuntimeError("missing admin_proof_b64")
            if not mgr.verify_admin(
                bytes.fromhex(target_rid_hex),
                _unb64(admin_proof_b64),
                transcript_hash,
            ):
                raise RuntimeError("admin authorization failed")

        if action == "server_info":
            return {"ok": True, "server_public_key_b64": server_pk_b64}

        if action == "create_registry":
            state_bytes = _unb64(cmd["state_bytes_b64"])
            state_cert = _unb64(cmd["state_cert_b64"])
            auth_rid = cmd.get("auth_registry_id")
            if not isinstance(auth_rid, str):
                raise RuntimeError("missing auth_registry_id")
            if not isinstance(admin_proof_b64, str) or not admin_proof_b64:
                raise RuntimeError("missing admin_proof_b64")
            tmp_mgr = zkac.RegistryManager()
            expected_rid = tmp_mgr.create(state_bytes, state_cert).hex()
            if expected_rid != auth_rid:
                raise RuntimeError("auth_registry_id does not match certified state")
            if not tmp_mgr.verify_admin(
                bytes.fromhex(expected_rid),
                _unb64(admin_proof_b64),
                transcript_hash,
            ):
                raise RuntimeError("admin authorization failed for create_registry")
            rid = mgr.create(state_bytes, state_cert)
            store.save_registry(rid.hex(), state_bytes, state_cert)
            return {"ok": True, "registry_id": rid.hex()}

        if action == "get_registry":
            rid_hex_cmd = cmd["registry_id"]
            _require_admin_for(rid_hex_cmd)
            rid = bytes.fromhex(rid_hex_cmd)
            state_bytes, state_cert = mgr.get(rid)
            return {
                "ok": True,
                "state_bytes_b64": _b64(state_bytes),
                "state_cert_b64": _b64(state_cert),
            }

        if action == "update_registry":
            rid_hex_cmd = cmd["registry_id"]
            _require_admin_for(rid_hex_cmd)
            rid = bytes.fromhex(rid_hex_cmd)
            state_bytes = _unb64(cmd["state_bytes_b64"])
            state_cert = _unb64(cmd["state_cert_b64"])
            mgr.update(rid, state_bytes, state_cert)
            store.save_registry(rid_hex_cmd, state_bytes, state_cert)
            return {"ok": True}

        return {"error": f"unknown command: {action}"}

    except Exception as exc:
        return {"error": str(exc)}


def _dispatch_fs(
    cmd: dict,
    store: _FileShareStore,
    ctx: dict,
) -> dict:
    """File-share commands authenticated via ``ctx`` (registry_id, role_id, issuance_pk_hex)."""
    try:
        action = cmd.get("cmd")
        registry_id_hex: str = ctx["registry_id_hex"]
        role_id: bytes = ctx["role_id"]
        issuance_pk_hex: str = ctx["issuance_pk_hex"]
        is_admin = role_id == zkac.admin_role_id()

        def _require_admin() -> None:
            if not is_admin:
                raise RuntimeError("admin role required for this command")

        if action == "whoami":
            return {
                "ok": True,
                "registry_id": registry_id_hex,
                "role_id": role_id.hex(),
                "is_admin": is_admin,
                "issuance_pk_hex": issuance_pk_hex,
            }

        if action == "bucket_create":
            _require_admin()
            bid = cmd.get("bucket_id") or uuid.uuid4().hex
            store.bucket_create(bid, registry_id_hex)
            return {"ok": True, "bucket_id": bid}

        if action == "bucket_put_blob":
            _require_admin()
            bid = cmd["bucket_id"]
            blob_id = cmd["blob_id"]
            ciphertext = _unb64(cmd["ciphertext_b64"])
            store.bucket_put_blob(bid, registry_id_hex, blob_id, ciphertext)
            return {"ok": True}

        if action == "bucket_put_role_grant":
            _require_admin()
            store.bucket_put_role_grant(
                cmd["bucket_id"],
                registry_id_hex,
                cmd["recipient_pk_hex"],
                cmd["role_id_hex"],
                int(cmd["acl_version"]),
                cmd["eph_pk_b64"],
                cmd["ciphertext_b64"],
            )
            return {"ok": True}

        if action == "bucket_set_role_acl":
            _require_admin()
            store.bucket_set_role_acl(
                cmd["bucket_id"],
                registry_id_hex,
                cmd["role_id_hex"],
                list(cmd.get("allowed_blob_ids", [])),
            )
            return {"ok": True}

        if action == "bucket_get_role_acl":
            _require_admin()
            role_acl = store.bucket_role_acl(cmd["bucket_id"], cmd["role_id_hex"])
            return {"ok": True, **role_acl}

        if action == "bucket_finalize":
            _require_admin()
            store.bucket_set_finalized(cmd["bucket_id"], registry_id_hex, True)
            return {"ok": True}

        if action == "bucket_delete":
            _require_admin()
            store.bucket_delete(cmd["bucket_id"], registry_id_hex)
            return {"ok": True}

        if action == "bucket_list_owned":
            _require_admin()
            return {"ok": True, "bucket_ids": store.buckets_owned_by(registry_id_hex)}

        if action == "fs_buckets":
            return {"ok": True, "bucket_ids": store.buckets_for_recipient(issuance_pk_hex)}

        if action == "fs_get_role_grant":
            bid = cmd["bucket_id"]
            grant = store.bucket_get_role_grant(bid, issuance_pk_hex)
            if not is_admin:
                expected_role = _safe_id(role_id.hex())
                granted_role_raw = grant.get("role_id_hex")
                if not isinstance(granted_role_raw, str) or not granted_role_raw:
                    raise RuntimeError("permissions updated; role grant is outdated, request a fresh grant")
                granted_role = _safe_id(granted_role_raw)
                if granted_role != expected_role:
                    raise RuntimeError("role grant not valid for authenticated role")
                current_acl = store.bucket_role_acl(bid, expected_role)
                acl_version = grant.get("acl_version")
                if not isinstance(acl_version, int):
                    raise RuntimeError("permissions updated; role grant is outdated, request a fresh grant")
                if acl_version != int(current_acl.get("version", -2)):
                    raise RuntimeError("permissions updated; request a fresh role grant")
            return {"ok": True, **grant}

        if action == "fs_get_blob":
            bid = cmd["bucket_id"]
            blob_id = cmd["blob_id"]
            if not is_admin and not store.bucket_blob_allowed_for_role(bid, role_id.hex(), blob_id):
                raise RuntimeError("blob access denied by role mask")
            data = store.bucket_get_blob(bid, blob_id)
            return {"ok": True, "ciphertext_b64": _b64(data)}

        return {"error": f"unknown command: {action}"}

    except Exception as exc:
        return {"error": str(exc)}


# ── per-connection handler ────────────────────────────────────────────

def _handle_conn(
    conn: socket.socket,
    addr: tuple,
    node: zkac.Node,
    mgr: zkac.RegistryManager,
    store: _FileShareStore,
    server_pk_b64: str,
    idle_timeout_s: float,
    slots: threading.BoundedSemaphore,
) -> None:
    peer = f"{addr[0]}:{addr[1]}"
    try:
        conn.settimeout(idle_timeout_s)
        session = server_handshake_anon(conn, node)
        framed = FramedSession(conn, session)
        transcript_hash = bytes(session.transcript_hash())

        hello = json.loads(framed.recv())
        op = hello.get("op")

        if op == "mgmt":
            while True:
                try:
                    cmd = json.loads(framed.recv())
                except (ConnectionError, OSError):
                    break
                resp = _dispatch_mgmt(cmd, mgr, store, server_pk_b64, transcript_hash)
                framed.send(json.dumps(resp).encode())
            return

        if op == "fs":
            try:
                registry_id = bytes.fromhex(hello["registry_id"])
                role_id = bytes.fromhex(hello["role_id"])
                proof = _unb64(hello["bbs_auth_b64"])
                issuance_pk_hex = hello["issuance_pk_hex"]
                _safe_id(issuance_pk_hex)
            except (KeyError, ValueError) as exc:
                framed.send(json.dumps({"error": f"invalid fs hello: {exc}"}).encode())
                return
            if role_id == zkac.admin_role_id():
                ok = mgr.verify_admin(registry_id, proof, transcript_hash)
            else:
                ok = mgr.verify_presentation(registry_id, role_id, proof, transcript_hash)
            if not ok:
                framed.send(json.dumps({"error": "auth failed"}).encode())
                return
            framed.send(json.dumps({"ok": True, "status": "authenticated"}).encode())
            ctx = {
                "registry_id_hex": registry_id.hex(),
                "role_id": role_id,
                "issuance_pk_hex": issuance_pk_hex,
            }
            while True:
                try:
                    cmd = json.loads(framed.recv())
                except (ConnectionError, OSError):
                    break
                resp = _dispatch_fs(cmd, store, ctx)
                framed.send(json.dumps(resp).encode())
            return

        framed.send(json.dumps({"error": f"unknown op: {op}"}).encode())

    except (ConnectionError, BrokenPipeError, OSError):
        pass
    except Exception as exc:
        print(f"[fs-server] {peer} error: {exc}")
        traceback.print_exc()
    finally:
        conn.close()
        slots.release()


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

def serve(
    data_dir: Path,
    host: str = "127.0.0.1",
    port: int = 9879,
    *,
    max_connections: int = 64,
    idle_timeout_s: float = 60.0,
    listen_backlog: int = 64,
    allow_non_loopback: bool = False,
) -> None:
    data_dir.mkdir(parents=True, exist_ok=True)
    store = _FileShareStore(data_dir)
    kp = store.load_or_create_keypair()
    server_pk_b64 = _b64(kp.public_key().to_bytes())
    pk_hex = kp.public_key().to_bytes().hex()
    node = zkac.Node(kp)

    mgr = zkac.RegistryManager()
    n = store.load_all_registries(mgr)

    print(f"[fs-server] data dir: {data_dir}")
    print(f"[fs-server] server transport public key (pin OOB): {pk_hex}")
    print(f"[fs-server] loaded {n} registries")
    print(f"[fs-server] listening on {host}:{port}")

    if not _is_loopback(host) and not allow_non_loopback:
        raise RuntimeError(
            "refusing to bind outside loopback. "
            "Use --allow-non-loopback only when exposure is intentional."
        )
    if not _is_loopback(host):
        print(f"[fs-server] warning: binding outside loopback: {host}:{port}", file=sys.stderr)

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((host, port))
    slots = threading.BoundedSemaphore(max_connections)
    sock.listen(listen_backlog)
    try:
        while True:
            conn, addr = sock.accept()
            if not slots.acquire(blocking=False):
                conn.close()
                continue
            threading.Thread(
                target=_handle_conn,
                args=(conn, addr, node, mgr, store, server_pk_b64, idle_timeout_s, slots),
                daemon=True,
            ).start()
    except KeyboardInterrupt:
        print("\n[fs-server] shutdown")
    finally:
        sock.close()


def main() -> None:
    p = argparse.ArgumentParser(description="ZKAC opaque file-share server (demo)")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=9879)
    p.add_argument(
        "--data-dir",
        type=Path,
        default=Path(__file__).resolve().parent / "fs_data",
        help="server state (transport key, registries, opaque buckets)",
    )
    p.add_argument("--idle-timeout", type=float, default=60.0)
    p.add_argument("--max-connections", type=int, default=64)
    p.add_argument("--allow-non-loopback", action="store_true")
    args = p.parse_args()

    serve(
        args.data_dir,
        host=args.host,
        port=args.port,
        max_connections=args.max_connections,
        idle_timeout_s=args.idle_timeout,
        allow_non_loopback=args.allow_non_loopback,
    )


if __name__ == "__main__":
    main()