Update app.py

app.py

@@ -1,7 +1,9 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
 from pydantic import BaseModel, Field
-from typing import Optional, List, Dict
+from typing import Optional, List, Dict, Any
 import httpx
 import json
 import logging
@@ -11,12 +13,16 @@ import hashlib
 from datetime import datetime, timezone
 import re
 import fastapi
-import glob
+import base64
+import imghdr
+import mimetypes
+import pathlib
+import uuid
 
 app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 
-# CORS (HF Spaces + static html)
+# Enable CORS for convenience (adjust origins for production)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -25,69 +31,70 @@ app.add_middleware(
 )
 
 # ---------------------------
-# 📦 Persistent storage
+# Persistent storage setup
 # ---------------------------
 PERSISTENT_ROOT = os.environ.get("PERSISTENT_DIR", "/data")
 if not os.path.isdir(PERSISTENT_ROOT):
+    # fallback to local data folder for dev
     PERSISTENT_ROOT = os.path.join(".", "data")
 os.makedirs(PERSISTENT_ROOT, exist_ok=True)
 
 CHAT_DIR = os.path.join(PERSISTENT_ROOT, "chat")
-DM_DIR = os.path.join(PERSISTENT_ROOT, "dm")
-USERS_DIR = os.path.join(PERSISTENT_ROOT, "users")
 os.makedirs(CHAT_DIR, exist_ok=True)
-os.makedirs(DM_DIR, exist_ok=True)
-os.makedirs(USERS_DIR, exist_ok=True)
 
+USERS_DIR = os.path.join(PERSISTENT_ROOT, "users")
+os.makedirs(USERS_DIR, exist_ok=True)
 USERS_FILE = os.path.join(USERS_DIR, "users.jsonl")
 SESSIONS_FILE = os.path.join(USERS_DIR, "sessions.json")
-ADMIN_KEY = os.environ.get("ADMIN_KEY", "")
+ROOMS_FILE = os.path.join(USERS_DIR, "rooms.json")   # rooms persisted here
+
+STATIC_DIR = os.path.join(PERSISTENT_ROOT, "static")
+IMAGES_DIR = os.path.join(STATIC_DIR, "images")
+os.makedirs(IMAGES_DIR, exist_ok=True)
+
+ADMIN_KEY = os.environ.get("ADMIN_KEY", "")  # set as secret in HF Space
+
+# Mount static files (serves uploaded images at /static/images/...)
+app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
-# Locks + ephemeral state
-app.state.file_locks: Dict[str, asyncio.Lock] = {}
+# Async locks
+app.state.chat_locks = {}
 app.state.users_lock = asyncio.Lock()
 app.state.sessions_lock = asyncio.Lock()
+app.state.rooms_lock = asyncio.Lock()
 app.state.sessions: Dict[str, str] = {}
 
-
-def _room_file(room_id: str) -> str:
-    h = hashlib.sha256(room_id.encode("utf-8")).hexdigest()[:32]
+# ---------------------------
+# Helpers: filenames, locks, time
+# ---------------------------
+def _chat_file_for(room_code: str) -> str:
+    h = hashlib.sha256(room_code.encode("utf-8")).hexdigest()[:32]
     return os.path.join(CHAT_DIR, f"{h}.jsonl")
 
-
-def _dm_file(a: str, b: str) -> str:
-    aa, bb = sorted([a, b])
-    k = hashlib.sha256(f"{aa}|{bb}".encode("utf-8")).hexdigest()[:40]
-    return os.path.join(DM_DIR, f"dm_{k}.jsonl")
-
-
 def _lock_for(path: str) -> asyncio.Lock:
-    lock = app.state.file_locks.get(path)
-    if not lock:
+    lock = app.state.chat_locks.get(path)
+    if lock is None:
         lock = asyncio.Lock()
-        app.state.file_locks[path] = lock
+        app.state.chat_locks[path] = lock
     return lock
 
-
 def _now_iso() -> str:
     return datetime.now(timezone.utc).isoformat()
 
+def _safe_handle(s: str) -> str:
+    s = (s or "").strip()
+    s = re.sub(r"[^a-zA-Z0-9_.~-]", "_", s)
+    return s[:24]
 
-def _valid_author(s: Optional[str]) -> str:
-    s = (s or "anon").strip()
-    s = re.sub(r"\s+", " ", s)
-    return s[:32] or "anon"
-
-
-def _valid_text(s: Optional[str]) -> str:
-    s = (s or "").rstrip("\n")
-    return s[:5000]
-
-
-def _is_data_url(img: str) -> bool:
-    return isinstance(img, str) and img.startswith("data:image/") and ";base64," in img
+def _hash_pin(pin: str) -> str:
+    return hashlib.sha256(("tlks|" + pin).encode("utf-8")).hexdigest()
 
+def _rand_token() -> str:
+    return hashlib.sha256(os.urandom(32)).hexdigest()
 
+# ---------------------------
+# JSONL read/write helpers
+# ---------------------------
 async def _append_jsonl(path: str, record: dict) -> None:
     line = json.dumps(record, ensure_ascii=False) + "\n"
     def _write():
@@ -95,7 +102,6 @@ async def _append_jsonl(path: str, record: dict) -> None:
             f.write(line)
     await asyncio.to_thread(_write)
 
-
 async def _read_jsonl(path: str) -> List[dict]:
     if not os.path.exists(path):
         return []
@@ -105,44 +111,25 @@ async def _read_jsonl(path: str) -> List[dict]:
     return await asyncio.to_thread(_read)
 
 # ---------------------------
-# 👤 Users & Auth helpers
+# Users & sessions
 # ---------------------------
-def _safe_handle(s: str) -> str:
-    s = (s or "").strip()
-    s = re.sub(r"[^a-zA-Z0-9_.~-]", "_", s)
-    return s[:24]
-
-def _hash_pin(pin: str) -> str:
-    return hashlib.sha256(("tlks|" + pin).encode("utf-8")).hexdigest()
-
-def _rand_token() -> str:
-    return hashlib.sha256(os.urandom(32)).hexdigest()
-
 async def _append_user(record: dict) -> None:
     line = json.dumps(record, ensure_ascii=False) + "\n"
-    def _write():
-        with open(USERS_FILE, "a", encoding="utf-8") as f:
-            f.write(line)
     async with app.state.users_lock:
+        def _write():
+            with open(USERS_FILE, "a", encoding="utf-8") as f:
+                f.write(line)
         await asyncio.to_thread(_write)
 
 async def _read_users() -> List[dict]:
     if not os.path.exists(USERS_FILE):
         return []
-    def _read():
-        with open(USERS_FILE, "r", encoding="utf-8") as f:
-            return [json.loads(x) for x in f if x.strip()]
     async with app.state.users_lock:
+        def _read():
+            with open(USERS_FILE, "r", encoding="utf-8") as f:
+                return [json.loads(x) for x in f if x.strip()]
         return await asyncio.to_thread(_read)
 
-async def _write_users(users: List[dict]):
-    async with app.state.users_lock:
-        def _write_all():
-            with open(USERS_FILE, "w", encoding="utf-8") as f:
-                for rec in users:
-                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")
-        await asyncio.to_thread(_write_all)
-
 async def _load_sessions():
     if os.path.exists(SESSIONS_FILE):
         def _read():
@@ -160,19 +147,106 @@ async def _save_sessions():
         await asyncio.to_thread(_write, app.state.sessions)
 
 @app.on_event("startup")
-async def _startup():
+async def _startup_load_sessions():
     await _load_sessions()
+    # ensure ROOMS_FILE exists with a default public 'main' room
+    async with app.state.rooms_lock:
+        if not os.path.exists(ROOMS_FILE):
+            default = [
+                {"code": "main", "name": "General", "private": False, "invite_code": "", "members": []}
+            ]
+            def _write():
+                with open(ROOMS_FILE, "w", encoding="utf-8") as f:
+                    json.dump(default, f)
+            await asyncio.to_thread(_write)
+
+# ---------------------------
+# Room persistence helpers
+# ---------------------------
+async def _read_rooms() -> List[dict]:
+    async with app.state.rooms_lock:
+        if not os.path.exists(ROOMS_FILE):
+            return []
+        def _read():
+            with open(ROOMS_FILE, "r", encoding="utf-8") as f:
+                return json.load(f)
+        return await asyncio.to_thread(_read)
+
+async def _save_rooms(rooms: List[dict]) -> None:
+    async with app.state.rooms_lock:
+        def _write(r):
+            with open(ROOMS_FILE, "w", encoding="utf-8") as f:
+                json.dump(r, f, ensure_ascii=False, indent=2)
+        await asyncio.to_thread(_write, rooms)
 
 def _require_admin(request: Request):
-    if not ADMIN_KEY or request.headers.get("x-admin-key") != ADMIN_KEY:
+    provided = request.headers.get("x-admin-key", "")
+    if not ADMIN_KEY or provided != ADMIN_KEY:
         raise fastapi.HTTPException(status_code=401, detail="Admin key required")
 
 # ---------------------------
-# 💬 Models
+# Image helpers (save data-URL or UploadFile)
+# ---------------------------
+def _guess_ext_from_mime(mime: str) -> str:
+    ext = mimetypes.guess_extension(mime.split(";")[0].strip()) or ""
+    if ext == ".jpe":
+        ext = ".jpg"
+    return ext
+
+def _safe_image_filename(data_bytes: bytes, ext: str) -> str:
+    h = hashlib.sha256(data_bytes).hexdigest()[:20]
+    filename = f"{h}{ext or ''}"
+    return filename
+
+async def save_data_url(data_url: str) -> str:
+    """
+    Accept data URL (data:<mime>;base64,<data>) and save to IMAGES_DIR.
+    Returns path relative to /static (e.g. /static/images/<file>).
+    """
+    if not data_url.startswith("data:"):
+        raise ValueError("Not a data URL")
+    header, b64 = data_url.split(",", 1)
+    # header e.g. "data:image/png;base64"
+    try:
+        payload = base64.b64decode(b64)
+    except Exception as e:
+        raise ValueError("Bad base64") from e
+
+    mime_match = header.split(";")[0].split(":", 1)[-1] if ";" in header else header.split(":", 1)[-1]
+    ext = _guess_ext_from_mime(mime_match) or (("." + imghdr.what(None, payload)) if imghdr.what(None, payload) else ".jpg")
+    fname = _safe_image_filename(payload, ext)
+    out_path = os.path.join(IMAGES_DIR, fname)
+    # save
+    def _write():
+        with open(out_path, "wb") as f:
+            f.write(payload)
+    await asyncio.to_thread(_write)
+    return f"/static/images/{fname}"
+
+async def save_upload_file(upload: UploadFile) -> str:
+    contents = await upload.read()
+    ext = ""
+    # try from content-type
+    if upload.content_type:
+        ext = _guess_ext_from_mime(upload.content_type)
+    if not ext:
+        # fallback to imghdr to detect
+        kind = imghdr.what(None, contents)
+        ext = f".{kind}" if kind else ".jpg"
+    fname = _safe_image_filename(contents, ext)
+    out_path = os.path.join(IMAGES_DIR, fname)
+    def _write():
+        with open(out_path, "wb") as f:
+            f.write(contents)
+    await asyncio.to_thread(_write)
+    return f"/static/images/{fname}"
+
+# ---------------------------
+# Models
 # ---------------------------
 class NewMessage(BaseModel):
     text: str = Field(..., min_length=1, max_length=5000)
-    images: Optional[List[str]] = Field(default=None, description="data:image/...;base64,...", max_items=4)
+    images: Optional[List[str]] = None  # can be data-URLs or already-servable URLs
 
 class NewUser(BaseModel):
     email: str
@@ -192,34 +266,14 @@ class UpdateUser(BaseModel):
     description: Optional[str] = None
     pin: Optional[str] = None
     disabled: Optional[bool] = None
-    following: Optional[List[str]] = None
-    followers: Optional[List[str]] = None
 
 class LoginReq(BaseModel):
     handle: str
     pin: str
 
-class MeProfilePatch(BaseModel):
-    profile_image: Optional[str] = None  # can be URL or data URL
-    description: Optional[str] = None
-
 # ---------------------------
-# 🔐 Auth
+# Auth utilities
 # ---------------------------
-@app.post("/auth/login")
-async def login(req: LoginReq):
-    users = await _read_users()
-    handle = _safe_handle(req.handle.lstrip("@"))
-    u = next((x for x in users if x.get("handle") == handle), None)
-    if not u or u.get("disabled"):
-        return {"ok": False, "error": "Invalid user"}
-    if u.get("pin_hash") != _hash_pin(req.pin):
-        return {"ok": False, "error": "Invalid PIN"}
-    token = _rand_token()
-    app.state.sessions[token] = handle
-    await _save_sessions()
-    return {"ok": True, "token": token}
-
 async def _require_user(request: Request) -> dict:
     auth = request.headers.get("authorization", "")
     if not auth.lower().startswith("bearer "):
@@ -234,28 +288,8 @@ async def _require_user(request: Request) -> dict:
         raise fastapi.HTTPException(status_code=401, detail="User disabled")
     return user
 
-@app.get("/me")
-async def me(request: Request):
-    user = await _require_user(request)
-    pub = {k: v for k, v in user.items() if k not in ("pin_hash",)}
-    return {"ok": True, "user": pub}
-
-@app.put("/me/profile")
-async def update_me_profile(patch: MeProfilePatch, request: Request):
-    me = await _require_user(request)
-    users = await _read_users()
-    for u in users:
-        if u.get("handle") == me["handle"]:
-            if patch.profile_image is not None:
-                u["profile_image"] = patch.profile_image.strip()
-            if patch.description is not None:
-                u["description"] = patch.description.strip()
-            break
-    await _write_users(users)
-    return {"ok": True}
-
 # ---------------------------
-# 👥 Admin Users
+# Admin: user creation (existing) and room management
 # ---------------------------
 @app.post("/admin/users")
 async def admin_create_user(user: NewUser, request: Request):
@@ -275,8 +309,8 @@ async def admin_create_user(user: NewUser, request: Request):
         "pin_hash": _hash_pin(user.pin),
         "disabled": False,
         "created_at": _now_iso(),
-        "following": [],
         "followers": [],
+        "following": [],
     }
     await _append_user(rec)
     pub = {k: v for k, v in rec.items() if k != "pin_hash"}
@@ -299,103 +333,214 @@ async def admin_update_user(handle: str, patch: UpdateUser, request: Request):
     _require_admin(request)
     handle = _safe_handle(handle)
     users = await _read_users()
-    found = False
+    changed = False
     for u in users:
         if u.get("handle") == handle:
-            data = patch.dict(exclude_unset=True)
-            if "pin" in data:
-                u["pin_hash"] = _hash_pin(data.pop("pin"))
-            for k, v in data.items():
-                u[k] = v
-            found = True
+            if patch.first_name is not None:
+                u["first_name"] = patch.first_name.strip()
+            if patch.last_name is not None:
+                u["last_name"] = patch.last_name.strip()
+            if patch.klass is not None:
+                u["klass"] = patch.klass.strip()
+            if patch.profile_image is not None:
+                u["profile_image"] = patch.profile_image.strip()
+            if patch.description is not None:
+                u["description"] = patch.description.strip()
+            if patch.disabled is not None:
+                u["disabled"] = bool(patch.disabled)
+            if patch.pin is not None:
+                u["pin_hash"] = _hash_pin(patch.pin)
+            changed = True
             break
-    if not found:
+    if not changed:
         raise fastapi.HTTPException(status_code=404, detail="User not found")
-    await _write_users(users)
+    # rewrite users file
+    async with app.state.users_lock:
+        def _write_all():
+            with open(USERS_FILE, "w", encoding="utf-8") as f:
+                for rec in users:
+                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+        await asyncio.to_thread(_write_all)
     return {"ok": True}
 
 # ---------------------------
-# 🌐 Public Users + Follow
+# Auth: login
 # ---------------------------
-@app.get("/users")
-async def public_users(q: Optional[str] = None, limit: int = 200):
-    users = await _read_users()
-    if q:
-        ql = q.lower()
-        users = [u for u in users if ql in u.get("handle", "").lower() or ql in u.get("first_name", "").lower() or ql in u.get("last_name", "").lower()]
-    users = users[: max(1, min(limit, 1000))]
-    out = []
-    for u in users:
-        if u.get("disabled"):
-            continue
-        out.append({
-            "handle": u["handle"],
-            "first_name": u["first_name"],
-            "last_name": u["last_name"],
-            "klass": u.get("klass", ""),
-            "profile_image": u.get("profile_image", ""),
-            "description": u.get("description", ""),
-            "followers": len(u.get("followers", [])),
-            "following": len(u.get("following", [])),
-        })
-    return {"ok": True, "users": out}
-
-@app.get("/user/{handle}")
-async def get_user_profile(handle: str):
+@app.post("/auth/login")
+async def login(req: LoginReq):
     users = await _read_users()
-    h = _safe_handle(handle)
-    u = next((x for x in users if x.get("handle") == h and not x.get("disabled")), None)
-    if not u:
-        raise fastapi.HTTPException(status_code=404, detail="User not found")
-    pub = {
-        "handle": u["handle"],
-        "first_name": u["first_name"],
-        "last_name": u["last_name"],
-        "klass": u.get("klass", ""),
-        "profile_image": u.get("profile_image", ""),
-        "description": u.get("description", ""),
-        "followers": u.get("followers", []),
-        "following": u.get("following", []),
-    }
+    handle = _safe_handle(req.handle.lstrip("@"))
+    u = next((x for x in users if x.get("handle") == handle), None)
+    if not u or u.get("disabled"):
+        return {"ok": False, "error": "Invalid user"}
+    if u.get("pin_hash") != _hash_pin(req.pin):
+        return {"ok": False, "error": "Invalid PIN"}
+    token = _rand_token()
+    app.state.sessions[token] = handle
+    await _save_sessions()
+    return {"ok": True, "token": token}
+
+@app.get("/me")
+async def me(request: Request):
+    user = await _require_user(request)
+    pub = {k: v for k, v in user.items() if k not in ("pin_hash",)}
     return {"ok": True, "user": pub}
 
-@app.post("/follow/{target}")
-async def follow_user(target: str, request: Request):
-    me = await _require_user(request)
-    target = _safe_handle(target)
-    users = await _read_users()
-    me_u = next((u for u in users if u.get("handle") == me["handle"]), None)
-    tgt_u = next((u for u in users if u.get("handle") == target and not u.get("disabled")), None)
-    if not tgt_u:
-        return {"ok": False, "error": "User not found"}
-    if target not in me_u.setdefault("following", []):
-        me_u["following"].append(target)
-    if me_u["handle"] not in tgt_u.setdefault("followers", []):
-        tgt_u["followers"].append(me_u["handle"])
-    await _write_users(users)
+# ---------------------------
+# Rooms endpoints
+# ---------------------------
+@app.get("/rooms")
+async def list_rooms(request: Request = None):
+    """
+    Returns rooms. If an Authorization bearer is supplied and the user is member/admin,
+    private room info will be visible. Otherwise invite_code is hidden.
+    """
+    rooms = await _read_rooms()
+    # If a valid bearer token present, get handle
+    authed_handle = None
+    auth = None
+    if request:
+        auth = request.headers.get("authorization", "")
+    if auth and auth.lower().startswith("bearer "):
+        token = auth.split(" ", 1)[1].strip()
+        authed_handle = app.state.sessions.get(token)
+    out = []
+    for r in rooms:
+        copy = {k: v for k, v in r.items() if k != "invite_code"}
+        # show invite_code only to admin (x-admin-key) or to admin bearer user (not implemented)
+        if request and request.headers.get("x-admin-key", "") == ADMIN_KEY:
+            copy["invite_code"] = r.get("invite_code", "")
+        # mark if current user is a member
+        copy["member"] = authed_handle in (r.get("members") or [])
+        out.append(copy)
+    return {"ok": True, "rooms": out}
+
+@app.post("/admin/rooms")
+async def admin_create_room(payload: dict, request: Request):
+    """
+    Admin-only: create room:
+    { "code": "grade6", "name": "Grade 6", "private": true }
+    Returns invite_code (server-generated) for private rooms.
+    """
+    _require_admin(request)
+    code = _safe_handle((payload.get("code") or "").lower())
+    if not code:
+        raise fastapi.HTTPException(status_code=400, detail="Invalid code")
+    name = (payload.get("name") or code).strip()
+    private = bool(payload.get("private", False))
+    rooms = await _read_rooms()
+    if any(r.get("code") == code for r in rooms):
+        return {"ok": False, "error": "Room exists"}
+    invite_code = ""
+    if private:
+        invite_code = hashlib.sha1(os.urandom(16)).hexdigest()[:10].upper()
+    rec = {"code": code, "name": name, "private": private, "invite_code": invite_code, "members": []}
+    rooms.append(rec)
+    await _save_rooms(rooms)
+    return {"ok": True, "room": rec}
+
+@app.get("/admin/rooms")
+async def admin_list_rooms(request: Request):
+    _require_admin(request)
+    rooms = await _read_rooms()
+    return {"ok": True, "rooms": rooms}
+
+@app.get("/admin/rooms/{code}")
+async def admin_get_room(code: str, request: Request):
+    _require_admin(request)
+    rooms = await _read_rooms()
+    r = next((x for x in rooms if x.get("code") == code), None)
+    if not r:
+        raise fastapi.HTTPException(status_code=404, detail="Room not found")
+    return {"ok": True, "room": r}
+
+@app.delete("/admin/rooms/{code}")
+async def admin_delete_room(code: str, request: Request):
+    _require_admin(request)
+    rooms = await _read_rooms()
+    code = _safe_handle(code)
+    new = [r for r in rooms if r.get("code") != code]
+    await _save_rooms(new)
     return {"ok": True}
 
-@app.post("/unfollow/{target}")
-async def unfollow_user(target: str, request: Request):
-    me = await _require_user(request)
-    target = _safe_handle(target)
-    users = await _read_users()
-    me_u = next((u for u in users if u.get("handle") == me["handle"]), None)
-    tgt_u = next((u for u in users if u.get("handle") == target), None)
-    if not tgt_u:
-        return {"ok": False, "error": "User not found"}
-    me_u["following"] = [h for h in me_u.get("following", []) if h != target]
-    tgt_u["followers"] = [h for h in tgt_u.get("followers", []) if h != me_u["handle"]]
-    await _write_users(users)
+@app.get("/admin/rooms/{code}/members")
+async def admin_room_members(code: str, request: Request):
+    _require_admin(request)
+    rooms = await _read_rooms()
+    r = next((x for x in rooms if x.get("code") == code), None)
+    if not r: raise fastapi.HTTPException(status_code=404, detail="Room not found")
+    return {"ok": True, "members": r.get("members", [])}
+
+@app.post("/admin/rooms/{code}/members/remove")
+async def admin_remove_room_member(code: str, body: dict, request: Request):
+    _require_admin(request)
+    handle = _safe_handle(body.get("handle", ""))
+    rooms = await _read_rooms()
+    found = False
+    for r in rooms:
+        if r.get("code") == code:
+            members = r.get("members", [])
+            if handle in members:
+                members.remove(handle)
+                r["members"] = members
+                found = True
+            break
+    if not found:
+        return {"ok": False, "error": "Member not found in room"}
+    await _save_rooms(rooms)
     return {"ok": True}
 
+@app.post("/rooms/{code}/join")
+async def join_room(code: str, body: dict, request: Request):
+    """
+    Authenticated users call to join private rooms using invite code.
+    Non-private rooms are joinable without invite (server adds user as member).
+    """
+    user = await _require_user(request)
+    code = _safe_handle(code)
+    rooms = await _read_rooms()
+    r = next((x for x in rooms if x.get("code") == code), None)
+    if not r:
+        raise fastapi.HTTPException(status_code=404, detail="Room not found")
+    if r.get("private"):
+        provided = (body.get("invite") or "").strip()
+        if not provided:
+            raise fastapi.HTTPException(status_code=400, detail="Invite required")
+        if provided != r.get("invite_code"):
+            raise fastapi.HTTPException(status_code=403, detail="Invalid invite code")
+    # add member if not already
+    members = set(r.get("members", []))
+    members.add(user["handle"])
+    r["members"] = list(members)
+    await _save_rooms(rooms)
+    return {"ok": True, "room": r}
+
 # ---------------------------
-# 💬 Rooms (main/public)
+# Chat endpoints (read + write) with privacy enforcement
 # ---------------------------
-@app.get("/chat/{room_id}")
-async def get_messages(room_id: str, limit: int = 50, since: Optional[str] = None):
+@app.get("/chat/{room_code}")
+async def get_messages(room_code: str, limit: int = 50, since: Optional[str] = None, request: Request = None):
+    """
+    Fetch messages for a room.
+    Public rooms are readable by anyone; private rooms require membership/auth.
+    """
     limit = max(1, min(limit, 200))
-    path = _room_file(room_id)
+    rooms = await _read_rooms()
+    room_code = _safe_handle(room_code)
+    room = next((r for r in rooms if r.get("code") == room_code), None)
+    if not room:
+        # allow fallback to legacy behavior if no room metadata: public
+        room = {"code": room_code, "private": False}
+    # if private, require membership
+    if room.get("private"):
+        # require bearer token
+        try:
+            user = await _require_user(request)
+        except Exception:
+            raise fastapi.HTTPException(status_code=401, detail="Authentication required for this room")
+        if user.get("handle") not in room.get("members", []):
+            raise fastapi.HTTPException(status_code=403, detail="Not a member of this room")
+    path = _chat_file_for(room_code)
     items = await _read_jsonl(path)
     if since:
         try:
@@ -409,126 +554,92 @@ async def get_messages(room_id: str, limit: int = 50, since: Optional[str] = Non
     items.sort(key=lambda m: m.get("created_at", ""))
     if len(items) > limit:
         items = items[-limit:]
-    return {"room_id": room_id, "count": len(items), "messages": items}
-
-@app.post("/chat/{room_id}")
-async def post_message(room_id: str, msg: NewMessage, request: Request):
+    return {"room": room_code, "count": len(items), "messages": items}
+
+@app.post("/chat/{room_code}")
+async def post_message(room_code: str, msg: NewMessage, request: Request):
+    """
+    Append a message to a room's chat.
+    Accepts images (data-urls or public URLs). Data-urls will be saved server-side and replaced with /static URLs.
+    """
+    room_code = _safe_handle(room_code)
+    rooms = await _read_rooms()
+    room = next((r for r in rooms if r.get("code") == room_code), None)
+    if not room:
+        # if room doesn't exist, disallow posting
+        raise fastapi.HTTPException(status_code=404, detail="Room not found")
+    # require auth
     user = await _require_user(request)
-    path = _room_file(room_id)
-    lock = _lock_for(path)
-
-    text = _valid_text(msg.text)
-    if not text:
+    if room.get("private") and user.get("handle") not in room.get("members", []):
+        raise fastapi.HTTPException(status_code=403, detail="Not a member of this room")
+    text = (msg.text or "").strip()
+    if not text and not msg.images:
         return {"ok": False, "error": "Empty message"}
-
-    images = msg.images or []
-    # accept only data URLs (already base64 from client)
-    images = [img for img in images if _is_data_url(img)][:4]
-
-    author = _valid_author((f"{user.get('first_name','')} {user.get('last_name','')}".strip()) or user["handle"])
+    images_out: List[str] = []
+    if msg.images:
+        # images can be data-urls or already public URLs; save data-urls
+        for im in msg.images[:8]:
+            try:
+                if isinstance(im, str) and im.startswith("data:"):
+                    url = await save_data_url(im)
+                    images_out.append(url)
+                else:
+                    # treat as already a served URL (validate basic)
+                    images_out.append(im)
+            except Exception as e:
+                logging.exception("image save failed")
+    author = f"{user.get('first_name','')} {user.get('last_name','')}".strip() or user.get("handle")
     created = _now_iso()
-    mid = hashlib.sha1(f"{room_id}|{author}|{created}|{text}".encode("utf-8")).hexdigest()[:16]
-    record = {
+    mid = hashlib.sha1(f"{room_code}|{author}|{created}|{text}".encode("utf-8")).hexdigest()[:16]
+    ip = request.client.host if request and request.client else None
+    rec = {
         "id": mid,
-        "room_id": room_id,
+        "room": room_code,
         "author": author,
-        "handle": user["handle"],
+        "handle": user.get("handle"),
         "klass": user.get("klass", ""),
         "profile_image": user.get("profile_image", ""),
         "text": text,
-        "images": images,
+        "images": images_out,
         "created_at": created,
+        "ip": ip,
+        "ua": request.headers.get("user-agent", "")[:200],
     }
-    async with lock:
-        await _append_jsonl(path, record)
-    return {"ok": True, "message": record}
-
-# ---------------------------
-# ✉️ DMs
-# ---------------------------
-@app.get("/dm/{target}")
-async def dm_get(target: str, request: Request, limit: int = 50, since: Optional[str] = None):
-    me = await _require_user(request)
-    target = _safe_handle(target)
-    users = await _read_users()
-    # You must follow to DM
-    me_u = next((u for u in users if u["handle"] == me["handle"]), None)
-    if target not in me_u.get("following", []):
-        raise fastapi.HTTPException(status_code=403, detail="Follow user first to DM")
-
-    path = _dm_file(me["handle"], target)
-    items = await _read_jsonl(path)
-    if since:
-        try:
-            since_dt = datetime.fromisoformat(since.replace("Z", "+00:00"))
-            items = [
-                m for m in items
-                if datetime.fromisoformat(str(m.get("created_at", "")).replace("Z", "+00:00")) > since_dt
-            ]
-        except Exception:
-            pass
-    items.sort(key=lambda m: m.get("created_at", ""))
-    if len(items) > limit:
-        items = items[-limit:]
-    return {"with": target, "count": len(items), "messages": items}
-
-@app.post("/dm/{target}")
-async def dm_post(target: str, msg: NewMessage, request: Request):
-    me = await _require_user(request)
-    target = _safe_handle(target)
-    users = await _read_users()
-    # must follow to DM
-    me_u = next((u for u in users if u["handle"] == me["handle"]), None)
-    if target not in me_u.get("following", []):
-        raise fastapi.HTTPException(status_code=403, detail="Follow user first to DM")
-    tgt_u = next((u for u in users if u.get("handle") == target and not u.get("disabled")), None)
-    if not tgt_u:
-        raise fastapi.HTTPException(status_code=404, detail="User not found")
-
-    text = _valid_text(msg.text)
-    if not text:
-        return {"ok": False, "error": "Empty message"}
-
-    images = [img for img in (msg.images or []) if _is_data_url(img)][:4]
-
-    path = _dm_file(me["handle"], target)
+    path = _chat_file_for(room_code)
     lock = _lock_for(path)
-
-    created = _now_iso()
-    mid = hashlib.sha1(f"{me['handle']}->{target}|{created}|{text}".encode("utf-8")).hexdigest()[:16]
-    rec = {
-        "id": mid,
-        "from": me["handle"],
-        "to": target,
-        "text": text,
-        "images": images,
-        "created_at": created,
-    }
     async with lock:
         await _append_jsonl(path, rec)
     return {"ok": True, "message": rec}
 
 # ---------------------------
-# 👤 Last N posts for a user (for profile popup)
+# Upload endpoint (multipart or JSON data-url)
 # ---------------------------
-def _iter_room_files() -> List[str]:
-    return glob.glob(os.path.join(CHAT_DIR, "*.jsonl"))
-
-@app.get("/user/{handle}/posts")
-async def user_posts(handle: str, limit: int = 5):
-    handle = _safe_handle(handle)
-    results: List[dict] = []
-    for path in _iter_room_files():
-        msgs = await _read_jsonl(path)
-        for m in msgs:
-            if m.get("handle") == handle:
-                results.append(m)
-    results.sort(key=lambda m: m.get("created_at", ""), reverse=True)
-    results = results[: max(1, min(limit, 20))]
-    return {"ok": True, "posts": results}
+@app.post("/upload_image")
+async def upload_image(file: Optional[UploadFile] = File(None), data_url: Optional[str] = Form(None), request: Request = None):
+    """
+    Accepts:
+    - multipart form 'file' (UploadFile)
+    - OR form field 'data_url' with a base64 data:... URL
+    Returns JSON: { ok: True, url: "/static/images/<file>" }
+    """
+    if file:
+        try:
+            url = await save_upload_file(file)
+            return {"ok": True, "url": url}
+        except Exception as e:
+            logging.exception("upload failed")
+            raise fastapi.HTTPException(status_code=400, detail="Upload failed")
+    if data_url:
+        try:
+            url = await save_data_url(data_url)
+            return {"ok": True, "url": url}
+        except Exception as e:
+            logging.exception("data-url save failed")
+            raise fastapi.HTTPException(status_code=400, detail="Bad data URL")
+    raise fastapi.HTTPException(status_code=400, detail="No file or data_url provided")
 
 # ---------------------------
-# (Original iCloud album endpoints) — unchanged
+# (leftover original icloud endpoints unchanged)
 # ---------------------------
 BASE_62_MAP = {c: i for i, c in enumerate("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")}
 
@@ -551,18 +662,28 @@ async def get_base_url(token: str) -> str:
         n = base62_to_int(token[1:3])
     return f"https://p{n:02d}-sharedstreams.icloud.com/{token}/sharedstreams/"
 
-ICLOUD_HEADERS = {"Origin": "https://www.icloud.com", "Content-Type": "text/plain"}
+ICLOUD_HEADERS = {
+    "Origin": "https://www.icloud.com",
+    "Content-Type": "text/plain",
+}
 ICLOUD_PAYLOAD = '{"streamCtag":null}'
 
 async def get_redirected_base_url(base_url: str, token: str) -> str:
     client = await get_client()
-    resp = await client.post(f"{base_url}webstream", headers=ICLOUD_HEADERS, data=ICLOUD_PAYLOAD, follow_redirects=False)
+    resp = await client.post(
+        f"{base_url}webstream", headers=ICLOUD_HEADERS, data=ICLOUD_PAYLOAD, follow_redirects=False
+    )
     if resp.status_code == 330:
-        body = resp.json()
-        host = body.get("X-Apple-MMe-Host")
-        if not host:
-            raise ValueError("Missing X-Apple-MMe-Host in 330 response")
-        return f"https://{host}/{token}/sharedstreams/"
+        try:
+            body = resp.json()
+            host = body.get("X-Apple-MMe-Host")
+            if not host:
+                raise ValueError("Missing X-Apple-MMe-Host in 330 response")
+            logging.info(f"Redirected to {host}")
+            return f"https://{host}/{token}/sharedstreams/"
+        except Exception as e:
+            logging.error(f"Redirect parsing failed: {e}")
+            raise
     elif resp.status_code == 200:
         return base_url
     else:
@@ -588,13 +709,16 @@ async def get_album(token: str):
     try:
         base_url = await get_base_url(token)
         base_url = await get_redirected_base_url(base_url, token)
+
         metadata = await get_metadata(base_url)
         guids = [photo["photoGuid"] for photo in metadata]
         asset_map = await get_asset_urls(base_url, guids)
+
         videos = []
         for photo in metadata:
             if photo.get("mediaAssetType", "").lower() != "video":
                 continue
+
             derivatives = photo.get("derivatives", {})
             best = max(
                 (d for k, d in derivatives.items() if k.lower() != "posterframe"),
@@ -603,19 +727,28 @@ async def get_album(token: str):
             )
             if not best:
                 continue
+
             checksum = best.get("checksum")
             info = asset_map.get(checksum)
             if not info:
                 continue
             video_url = f"https://{info['url_location']}{info['url_path']}"
-            poster = ""
+
+            poster = None
             pf = derivatives.get("PosterFrame")
             if pf:
                 pf_info = asset_map.get(pf.get("checksum"))
                 if pf_info:
                     poster = f"https://{pf_info['url_location']}{pf_info['url_path']}"
-            videos.append({"caption": photo.get("caption", ""), "url": video_url, "poster": poster})
+
+            videos.append({
+                "caption": photo.get("caption", ""),
+                "url": video_url,
+                "poster": poster or "",
+            })
+
         return {"videos": videos}
+
     except Exception as e:
         logging.exception("Error in get_album")
         return {"error": str(e)}