Update app.py

app.py

@@ -11,11 +11,12 @@ import hashlib
 from datetime import datetime, timezone
 import re
 import fastapi
+import glob
 
 app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 
-# Enable CORS for all origins
+# CORS (HF Spaces + static html)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -24,38 +25,47 @@ app.add_middleware(
 )
 
 # ---------------------------
-# 📦 Persistent storage setup
+# 📦 Persistent storage
 # ---------------------------
 PERSISTENT_ROOT = os.environ.get("PERSISTENT_DIR", "/data")
-CHAT_DIR = os.path.join(PERSISTENT_ROOT, "chat")
 if not os.path.isdir(PERSISTENT_ROOT):
-    CHAT_DIR = os.path.join(".", "data", "chat")
-os.makedirs(CHAT_DIR, exist_ok=True)
+    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_FILE = os.path.join(USERS_DIR, "users.jsonl")
 SESSIONS_FILE = os.path.join(USERS_DIR, "sessions.json")
-ADMIN_KEY = os.environ.get("ADMIN_KEY", "")  # set in HF Space secrets
+ADMIN_KEY = os.environ.get("ADMIN_KEY", "")
 
-# Async file locks
-app.state.chat_locks = {}
+# Locks + ephemeral state
+app.state.file_locks: Dict[str, asyncio.Lock] = {}
 app.state.users_lock = asyncio.Lock()
 app.state.sessions_lock = asyncio.Lock()
 app.state.sessions: Dict[str, str] = {}
 
 
-def _chat_file_for(video_id: str) -> str:
-    # Stable filename to avoid path traversal
-    h = hashlib.sha256(video_id.encode("utf-8")).hexdigest()[:32]
+def _room_file(room_id: str) -> str:
+    h = hashlib.sha256(room_id.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.chat_locks.get(path)
-    if lock is None:
+    lock = app.state.file_locks.get(path)
+    if not lock:
         lock = asyncio.Lock()
-        app.state.chat_locks[path] = lock
+        app.state.file_locks[path] = lock
     return lock
 
 
@@ -69,72 +79,69 @@ def _valid_author(s: Optional[str]) -> str:
     return s[:32] or "anon"
 
 
-def _valid_text(s: str) -> str:
+def _valid_text(s: Optional[str]) -> str:
     s = (s or "").rstrip("\n")
-    return s[:2000]
+    return s[:5000]
+
+
+def _is_data_url(img: str) -> bool:
+    return isinstance(img, str) and img.startswith("data:image/") and ";base64," in img
 
 
 async def _append_jsonl(path: str, record: dict) -> None:
     line = json.dumps(record, ensure_ascii=False) + "\n"
-
     def _write():
         with open(path, "a", encoding="utf-8") as f:
             f.write(line)
-
     await asyncio.to_thread(_write)
 
 
 async def _read_jsonl(path: str) -> List[dict]:
     if not os.path.exists(path):
         return []
-
     def _read():
         with open(path, "r", encoding="utf-8") as f:
             return [json.loads(x) for x in f if x.strip()]
-
     return await asyncio.to_thread(_read)
 
-
 # ---------------------------
 # 👤 Users & Auth helpers
 # ---------------------------
-
 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:
         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:
         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):
@@ -145,66 +152,28 @@ async def _load_sessions():
     else:
         app.state.sessions = {}
 
-
 async def _save_sessions():
     def _write(data):
         with open(SESSIONS_FILE, "w", encoding="utf-8") as f:
             json.dump(data, f)
-
     async with app.state.sessions_lock:
         await asyncio.to_thread(_write, app.state.sessions)
 
-
 @app.on_event("startup")
-async def _startup_load_sessions():
+async def _startup():
     await _load_sessions()
 
-
 def _require_admin(request: Request):
-    provided = request.headers.get("x-admin-key", "")
-    if not ADMIN_KEY or provided != ADMIN_KEY:
+    if not ADMIN_KEY or request.headers.get("x-admin-key") != ADMIN_KEY:
         raise fastapi.HTTPException(status_code=401, detail="Admin key required")
 
-
 # ---------------------------
-# 💬 Chat API (public read, authed write)
+# 💬 Models
 # ---------------------------
 class NewMessage(BaseModel):
-    author: Optional[str] = Field(default="anon", max_length=64)  # ignored on server write
     text: str = Field(..., min_length=1, max_length=5000)
+    images: Optional[List[str]] = Field(default=None, description="data:image/...;base64,...", max_items=4)
 
-
-@app.get("/chat/{video_id}")
-async def get_messages(video_id: str, limit: int = 50, since: Optional[str] = None):
-    """
-    Fetch messages for a room/video.
-    - limit: max messages (default 50)
-    - since: ISO8601 timestamp; return only messages newer than this
-    """
-    limit = max(1, min(limit, 200))
-    path = _chat_file_for(video_id)
-    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  # Ignore bad since param
-
-    items.sort(key=lambda m: m.get("created_at", ""))
-    if len(items) > limit:
-        items = items[-limit:]
-
-    return {"video_id": video_id, "count": len(items), "messages": items}
-
-
-# ---------------------------
-# 👤 Users & Auth API
-# ---------------------------
 class NewUser(BaseModel):
     email: str
     first_name: str
@@ -215,7 +184,6 @@ class NewUser(BaseModel):
     description: Optional[str] = ""
     pin: str = Field(..., min_length=3, max_length=32)
 
-
 class UpdateUser(BaseModel):
     first_name: Optional[str] = None
     last_name: Optional[str] = None
@@ -224,22 +192,78 @@ 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
+# ---------------------------
+@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 "):
+        raise fastapi.HTTPException(status_code=401, detail="Missing bearer token")
+    token = auth.split(" ", 1)[1].strip()
+    handle = app.state.sessions.get(token)
+    if not handle:
+        raise fastapi.HTTPException(status_code=401, detail="Invalid session")
+    users = await _read_users()
+    user = next((u for u in users if u.get("handle") == handle), None)
+    if not user or user.get("disabled"):
+        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
+# ---------------------------
 @app.post("/admin/users")
 async def admin_create_user(user: NewUser, request: Request):
     _require_admin(request)
     users = await _read_users()
-
     handle = _safe_handle(user.handle.lstrip("@"))
     if any(u.get("handle") == handle for u in users):
         return {"ok": False, "error": "Handle already exists"}
-
     rec = {
         "email": user.email.strip(),
         "first_name": user.first_name.strip(),
@@ -251,12 +275,13 @@ async def admin_create_user(user: NewUser, request: Request):
         "pin_hash": _hash_pin(user.pin),
         "disabled": False,
         "created_at": _now_iso(),
+        "following": [],
+        "followers": [],
     }
     await _append_user(rec)
     pub = {k: v for k, v in rec.items() if k != "pin_hash"}
     return {"ok": True, "user": pub}
 
-
 @app.get("/admin/users")
 async def admin_list_users(request: Request, q: Optional[str] = None, limit: int = 200):
     _require_admin(request)
@@ -269,145 +294,255 @@ async def admin_list_users(request: Request, q: Optional[str] = None, limit: int
         u.pop("pin_hash", None)
     return {"count": len(users), "users": users}
 
-
 @app.put("/admin/users/{handle}")
 async def admin_update_user(handle: str, patch: UpdateUser, request: Request):
     _require_admin(request)
     handle = _safe_handle(handle)
     users = await _read_users()
-    changed = False
+    found = False
     for u in users:
         if u.get("handle") == handle:
-            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
+            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
             break
-    if not changed:
+    if not found:
         raise fastapi.HTTPException(status_code=404, detail="User not found")
-
-    # rewrite 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)
+    await _write_users(users)
     return {"ok": True}
 
-
-@app.post("/auth/login")
-async def login(req: LoginReq):
+# ---------------------------
+# 🌐 Public Users + Follow
+# ---------------------------
+@app.get("/users")
+async def public_users(q: Optional[str] = None, limit: int = 200):
     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}
-
+    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):
+    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", []),
+    }
+    return {"ok": True, "user": pub}
 
-async def _require_user(request: Request) -> dict:
-    auth = request.headers.get("authorization", "")
-    if not auth.lower().startswith("bearer "):
-        raise fastapi.HTTPException(status_code=401, detail="Missing bearer token")
-    token = auth.split(" ", 1)[1].strip()
-    handle = app.state.sessions.get(token)
-    if not handle:
-        raise fastapi.HTTPException(status_code=401, detail="Invalid session")
+@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)
+    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()
-    user = next((u for u in users if u.get("handle") == handle), None)
-    if not user or user.get("disabled"):
-        raise fastapi.HTTPException(status_code=401, detail="User disabled")
-    return user
+    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)
+    return {"ok": True}
 
+# ---------------------------
+# 💬 Rooms (main/public)
+# ---------------------------
+@app.get("/chat/{room_id}")
+async def get_messages(room_id: str, limit: int = 50, since: Optional[str] = None):
+    limit = max(1, min(limit, 200))
+    path = _room_file(room_id)
+    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 {"room_id": room_id, "count": len(items), "messages": items}
 
-@app.get("/me")
-async def me(request: Request):
+@app.post("/chat/{room_id}")
+async def post_message(room_id: str, msg: NewMessage, 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("/chat/{video_id}")
-async def post_message(video_id: str, msg: NewMessage, request: Request):
-    """
-    Append a message to a room's chat.
-    Body: { "text": "hello" }
-    """
-    path = _chat_file_for(video_id)
+    path = _room_file(room_id)
     lock = _lock_for(path)
 
-    # Require a valid session and take author from the user profile
-    user = await _require_user(request)
-    author = _valid_author(f"{user.get('first_name','')} {user.get('last_name','')}".strip() or user["handle"])
-    handle = user["handle"]
-    klass = user.get("klass", "")
-    profile_image = user.get("profile_image", "")
-
     text = _valid_text(msg.text)
     if not text:
         return {"ok": False, "error": "Empty message"}
 
-    created = _now_iso()
-    mid = hashlib.sha1(f"{video_id}|{author}|{created}|{text}".encode("utf-8")).hexdigest()[:16]
-    ip = request.client.host if request and request.client else None
+    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"])
+    created = _now_iso()
+    mid = hashlib.sha1(f"{room_id}|{author}|{created}|{text}".encode("utf-8")).hexdigest()[:16]
     record = {
         "id": mid,
-        "video_id": video_id,
+        "room_id": room_id,
         "author": author,
-        "handle": handle,
-        "klass": klass,
-        "profile_image": profile_image,
+        "handle": user["handle"],
+        "klass": user.get("klass", ""),
+        "profile_image": user.get("profile_image", ""),
         "text": text,
+        "images": images,
         "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)
+    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}
 
 # ---------------------------
-# (Original iCloud album endpoints)
+# 👤 Last N posts for a user (for profile popup)
 # ---------------------------
-BASE_62_MAP = {c: i for i, c in enumerate("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")}
+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}
+
+# ---------------------------
+# (Original iCloud album endpoints) — unchanged
+# ---------------------------
+BASE_62_MAP = {c: i for i, c in enumerate("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")}
 
 async def get_client() -> httpx.AsyncClient:
     if not hasattr(app.state, "client"):
         app.state.client = httpx.AsyncClient(timeout=15.0)
     return app.state.client
 
-
 def base62_to_int(token: str) -> int:
     result = 0
     for ch in token:
         result = result * 62 + BASE_62_MAP[ch]
     return result
 
-
 async def get_base_url(token: str) -> str:
     first = token[0]
     if first == "A":
@@ -416,69 +551,50 @@ 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:
-        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
+        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/"
     elif resp.status_code == 200:
         return base_url
     else:
         resp.raise_for_status()
 
-
 async def post_json(path: str, base_url: str, payload: str) -> dict:
     client = await get_client()
     resp = await client.post(f"{base_url}{path}", headers=ICLOUD_HEADERS, data=payload)
     resp.raise_for_status()
     return resp.json()
 
-
 async def get_metadata(base_url: str) -> list:
     data = await post_json("webstream", base_url, ICLOUD_PAYLOAD)
     return data.get("photos", [])
 
-
 async def get_asset_urls(base_url: str, guids: list) -> dict:
     payload = json.dumps({"photoGuids": guids})
     data = await post_json("webasseturls", base_url, payload)
     return data.get("items", {})
 
-
 @app.get("/album/{token}")
 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"),
@@ -487,33 +603,23 @@ 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 = None
+            poster = ""
             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 or "",
-            })
-
+            videos.append({"caption": photo.get("caption", ""), "url": video_url, "poster": poster})
         return {"videos": videos}
-
     except Exception as e:
         logging.exception("Error in get_album")
         return {"error": str(e)}
 
-
 @app.get("/album/{token}/raw")
 async def get_album_raw(token: str):
     try:
@@ -525,4 +631,4 @@ async def get_album_raw(token: str):
         return {"metadata": metadata, "asset_urls": asset_map}
     except Exception as e:
         logging.exception("Error in get_album_raw")
-        return {"error": str(e)}
+        return {"error": str(e)}