"""
MeshBay Hub POC v1
Minimal FastAPI hub: user registration, JWT issuance, node announcement.
In-memory storage only — not persistent across restarts.
"""

from fastapi import FastAPI, HTTPException, Depends, Header
from pydantic import BaseModel
from cryptography.hazmat.primitives.kdf.argon2 import Argon2id
import jwt, uuid, os, time, base64

app = FastAPI(title="MeshBay Hub POC", version="0.1.0")

HUB_ID = "meshbay.org"
MNP_VERSION = "0.1"
MHP_VERSION = "0.1"
ACCESS_TOKEN_TTL = 3600        # 1 hour
REFRESH_TOKEN_TTL = 86400 * 30  # 30 days

# Load hub keypair (generated once with gen_hub_keys.py)
with open("hub_private.pem", "rb") as f:
    HUB_SK_PEM = f.read()
with open("hub_public.pem", "rb") as f:
    HUB_PK_PEM = f.read()

# In-memory stores (POC — lost on restart)
users: dict = {}         # username → user record
nodes: dict = {}         # node_id → node record
refresh_tokens: dict = {}  # token → user_id
groups: dict = {}        # group_id → group record
gek_bundles: dict = {}   # (group_id, user_id) → encrypted GEK bundle


# ── Models ────────────────────────────────────────────────────────────────────

class UserRegister(BaseModel):
    username: str
    password: str
    pk_user_ed25519: str  # base64 raw 32 bytes
    pk_user_x25519: str   # base64 raw 32 bytes

class UserLogin(BaseModel):
    username: str
    password: str

class RefreshRequest(BaseModel):
    refresh_token: str

class NodeAnnounce(BaseModel):
    pk_node: str                    # base64 Ed25519 raw public key
    endpoint_hint: str | None = None  # "ip:port" discovered via STUN/UPnP

class GroupCreate(BaseModel):
    name: str

class GEKBundle(BaseModel):
    pk_eph_b64:  str   # ephemeral X25519 public key used during wrapping
    nonce_b64:   str   # ChaCha20-Poly1305 nonce
    wrapped_b64: str   # encrypted GEK (opaque to hub)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _hash_password(password: str) -> tuple[bytes, bytes]:
    salt = os.urandom(16)
    kdf = Argon2id(salt=salt, length=32, iterations=3, lanes=4, memory_cost=65536)
    return kdf.derive(password.encode()), salt

def _verify_password(password: str, pw_hash: bytes, salt: bytes) -> bool:
    try:
        Argon2id(salt=salt, length=32, iterations=3, lanes=4, memory_cost=65536
                 ).verify(password.encode(), pw_hash)
        return True
    except Exception:
        return False

def _issue_access_token(user: dict) -> str:
    now = int(time.time())
    payload = {
        "iss": HUB_ID,
        "sub": user["user_id"],
        "pk_user": user["pk_ed25519"],
        "hub_id": HUB_ID,
        "jti": str(uuid.uuid4()),   # unique per token — enables revocation, prevents replay
        "iat": now,
        "exp": now + ACCESS_TOKEN_TTL,
    }
    return jwt.encode(payload, HUB_SK_PEM, algorithm="EdDSA")

def _get_current_user(authorization: str = Header(...)) -> dict:
    try:
        scheme, token = authorization.split(None, 1)
        if scheme.lower() != "bearer":
            raise ValueError("Not bearer")
        payload = jwt.decode(token, HUB_PK_PEM, algorithms=["EdDSA"])
        user = next((u for u in users.values() if u["user_id"] == payload["sub"]), None)
        if not user:
            raise HTTPException(status_code=401, detail="User not found")
        return user
    except HTTPException:
        raise
    except Exception:
        raise HTTPException(status_code=401, detail="Invalid or expired token")


# ── Endpoints ─────────────────────────────────────────────────────────────────

@app.get("/v1/hub/info")
def hub_info():
    return {
        "hub_id": HUB_ID,
        "mnp_version": MNP_VERSION,
        "mhp_version": MHP_VERSION,
        "users": len(users),
        "nodes": len(nodes),
    }

@app.get("/v1/hub/pubkey")
def hub_pubkey():
    """Return hub Ed25519 public key PEM — nodes cache this on first contact."""
    return {"pk_hub_pem": HUB_PK_PEM.decode()}

@app.post("/v1/users/register", status_code=201)
def register(body: UserRegister):
    if body.username in users:
        raise HTTPException(status_code=409, detail="Username already taken")
    if len(body.password) < 8:
        raise HTTPException(status_code=422, detail="Password too short")
    pw_hash, pw_salt = _hash_password(body.password)
    user_id = str(uuid.uuid4())
    users[body.username] = {
        "user_id": user_id,
        "username": body.username,
        "pw_hash": pw_hash,
        "pw_salt": pw_salt,
        "pk_ed25519": body.pk_user_ed25519,
        "pk_x25519":  body.pk_user_x25519,
        "created_at": int(time.time()),
    }
    return {"user_id": user_id}

@app.post("/v1/users/login")
def login(body: UserLogin):
    user = users.get(body.username)
    if not user or not _verify_password(body.password, user["pw_hash"], user["pw_salt"]):
        raise HTTPException(status_code=401, detail="Invalid credentials")
    access_token = _issue_access_token(user)
    refresh_token = base64.urlsafe_b64encode(os.urandom(32)).decode()
    refresh_tokens[refresh_token] = user["user_id"]
    return {
        "access_token": access_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
        "expires_in": ACCESS_TOKEN_TTL,
    }

@app.post("/v1/users/token/refresh")
def token_refresh(body: RefreshRequest):
    user_id = refresh_tokens.get(body.refresh_token)
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid refresh token")
    user = next((u for u in users.values() if u["user_id"] == user_id), None)
    if not user:
        raise HTTPException(status_code=401, detail="User not found")
    return {
        "access_token": _issue_access_token(user),
        "token_type": "bearer",
        "expires_in": ACCESS_TOKEN_TTL,
    }

@app.post("/v1/nodes/announce", status_code=201)
def announce_node(body: NodeAnnounce, user: dict = Depends(_get_current_user)):
    node_id = str(uuid.uuid4())
    nodes[node_id] = {
        "node_id": node_id,
        "user_id": user["user_id"],
        "username": user["username"],
        "pk_node": body.pk_node,
        "endpoint_hint": body.endpoint_hint,
        "announced_at": int(time.time()),
    }
    return {"node_id": node_id}

@app.get("/v1/nodes/{node_id}")
def get_node(node_id: str, user: dict = Depends(_get_current_user)):
    node = nodes.get(node_id)
    if not node:
        raise HTTPException(status_code=404, detail="Node not found")
    return {
        "node_id": node["node_id"],
        "username": node["username"],
        "pk_node": node["pk_node"],
        "endpoint_hint": node["endpoint_hint"],
        "announced_at": node["announced_at"],
    }

@app.get("/v1/users/{username}/pubkeys")
def get_user_pubkeys(username: str, user: dict = Depends(_get_current_user)):
    """Return a user's public keys so the admin can wrap the GEK for them."""
    target = users.get(username)
    if not target:
        raise HTTPException(status_code=404, detail="User not found")
    return {
        "user_id":    target["user_id"],
        "username":   target["username"],
        "pk_ed25519": target["pk_ed25519"],
        "pk_x25519":  target["pk_x25519"],
    }

@app.post("/v1/groups", status_code=201)
def create_group(body: GroupCreate, user: dict = Depends(_get_current_user)):
    group_id = str(uuid.uuid4())
    groups[group_id] = {
        "group_id":   group_id,
        "name":       body.name,
        "admin_id":   user["user_id"],
        "admin_name": user["username"],
        "created_at": int(time.time()),
        "members":    [user["user_id"]],
    }
    return {"group_id": group_id, "name": body.name}

@app.post("/v1/groups/{group_id}/members/{username}/gek", status_code=201)
def store_gek_bundle(
    group_id: str, username: str,
    body: GEKBundle,
    user: dict = Depends(_get_current_user)
):
    """Admin stores an encrypted GEK bundle for a group member.
    The hub stores the bundle opaquely — it cannot decrypt it."""
    group = groups.get(group_id)
    if not group:
        raise HTTPException(status_code=404, detail="Group not found")
    if group["admin_id"] != user["user_id"]:
        raise HTTPException(status_code=403, detail="Only group admin can add members")
    target = users.get(username)
    if not target:
        raise HTTPException(status_code=404, detail="User not found")

    key = (group_id, target["user_id"])
    gek_bundles[key] = {
        "group_id":   group_id,
        "user_id":    target["user_id"],
        "pk_eph_b64":  body.pk_eph_b64,
        "nonce_b64":   body.nonce_b64,
        "wrapped_b64": body.wrapped_b64,
        "stored_at":  int(time.time()),
    }
    if target["user_id"] not in group["members"]:
        group["members"].append(target["user_id"])
    return {"status": "stored", "group_id": group_id, "username": username}

@app.get("/v1/groups/{group_id}/gek")
def get_my_gek_bundle(group_id: str, user: dict = Depends(_get_current_user)):
    """Authenticated member retrieves their own encrypted GEK bundle."""
    group = groups.get(group_id)
    if not group:
        raise HTTPException(status_code=404, detail="Group not found")
    key = (group_id, user["user_id"])
    bundle = gek_bundles.get(key)
    if not bundle:
        raise HTTPException(status_code=404, detail="No GEK bundle for this user in this group")
    return {
        "group_id":   group_id,
        "pk_eph_b64":  bundle["pk_eph_b64"],
        "nonce_b64":   bundle["nonce_b64"],
        "wrapped_b64": bundle["wrapped_b64"],
    }