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

import base64
import os
import time
import uuid

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

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"],
    }