"""Persistent state for RigShare: accounts, sessions, guests, config, rooms
and snapshots.

Everything lives as small JSON files in ``<ComfyUI>/rigshare``. Writes are atomic
(write to a temp file then rename) so a crash never leaves a half-written file.
"""

import hashlib
import hmac
import json
import logging
import os
import re
import secrets
import threading
import time

from . import jwt

log = logging.getLogger("ComfyUI-RigShare")

PBKDF2_ITERATIONS = 240_000
SESSION_TTL = 30 * 24 * 3600      # a login lasts at most this long
TOKEN_TTL = 12 * 3600              # each JWT is short-lived...
TOKEN_REFRESH = 6 * 3600           # ...and re-issued while in use once older than this
ISSUER = "rigshare"
USERNAME_RE = re.compile(r"^(?!(?i:con|prn|aux|nul|com\d|lpt\d)(?:\.|$))[A-Za-z0-9_][A-Za-z0-9_.-]{0,30}[A-Za-z0-9_-]\Z")  # no leading/trailing dot, no Windows device names
PERMS = ("edit", "queue", "manager", "admin")
API_KEY_PREFIX = "rs_"

DEFAULT_CONFIG = {
    "server_name": "RigShare",
    # Everyone must log in before ComfyUI loads (login page + session cookie).
    # API clients authenticate with a personal API key.
    "require_login": True,
    # Only when require_login is off: visitors without an account can watch and chat.
    "allow_guests": False,
    "guest_permissions": {"edit": False, "queue": False},
    # Relay execution progress/previews to every connected browser, not only
    # the one that queued the prompt.
    "broadcast_execution": True,
    # Reject queue/interrupt/upload HTTP calls from users lacking permission.
    # Requests from trusted IPs skip both this and the login requirement.
    "api_protection": {
        "enabled": True,
        "trusted_ips": ["127.0.0.1", "::1"],
    },
    # Hide ComfyUI's own Comfy.org account sign-in (top-bar button and dialog).
    "hide_comfy_account": False,
    # ComfyUI's own Workflows and Apps sidebar tabs duplicate RigShare's Files tab.
    "hide_comfy_file_tabs": True,
    # ComfyUI Manager in its own sidebar tab instead of the top-bar button.
    "manager_tab": False,
    # Remove templates that need paid API nodes from the template browser.
    "hide_api_templates": False,
    # Rooms nobody opened for this many days are forgotten.
    "room_expiry_days": 14,
    "snapshot_interval_sec": 300,
    "snapshot_keep": 30,
    "chat_keep": 0,  # 0 keeps the whole chat; N trims it to the last N messages at startup
}


def _hash_password(password, salt=None):
    salt = salt or secrets.token_bytes(16)
    digest = hashlib.pbkdf2_hmac("sha256", password.encode(), salt, PBKDF2_ITERATIONS)
    return f"pbkdf2_sha256${PBKDF2_ITERATIONS}${salt.hex()}${digest.hex()}"


def _verify_password(password, stored):
    try:
        _, iterations, salt, digest = stored.split("$")
        candidate = hashlib.pbkdf2_hmac("sha256", password.encode(), bytes.fromhex(salt), int(iterations))
        return hmac.compare_digest(candidate.hex(), digest)
    except (ValueError, AttributeError):
        return False


def _deep_merge(base, override):
    out = dict(base)
    for key, value in (override or {}).items():
        if isinstance(value, dict) and isinstance(out.get(key), dict):
            out[key] = _deep_merge(out[key], value)
        else:
            out[key] = value
    return out


class Store:
    def __init__(self, folder):
        self.folder = folder
        self.snapshot_dir = os.path.join(folder, "snapshots")
        self.room_dir = os.path.join(folder, "rooms")
        os.makedirs(self.snapshot_dir, exist_ok=True)
        os.makedirs(self.room_dir, exist_ok=True)
        self._lock = threading.RLock()

        self.secret = self._load_secret()
        self.config = _deep_merge(DEFAULT_CONFIG, self._read("config.json", {}))
        self._write("config.json", self.config)
        self.users = self._read("users.json", {})
        self.sessions = self._read("sessions.json", {})
        self.guests = self._read("guests.json", {})

        self._provision_admin()
        self._prune_sessions()

    # ----- file helpers -------------------------------------------------

    def _path(self, name):
        return os.path.join(self.folder, name)

    def _read(self, name, default):
        try:
            with open(self._path(name), "r", encoding="utf-8") as f:
                return json.load(f)
        except FileNotFoundError:
            return default
        except (OSError, ValueError) as e:
            log.error(f"[RigShare] Could not read {name}: {e}")
            return default

    def _write(self, name, data):
        with self._lock:
            path = self._path(name)
            tmp = f"{path}.tmp"
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)
            os.replace(tmp, path)

    # ----- setup --------------------------------------------------------

    def _load_secret(self):
        """JWT signing key, generated once and kept private to the server."""
        path = self._path("secret.key")
        try:
            with open(path, "rb") as f:
                secret = f.read()
            if len(secret) >= 32:
                return secret
        except FileNotFoundError:
            pass
        secret = secrets.token_bytes(64)
        fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(fd, "wb") as f:
            f.write(secret)
        return secret

    def _provision_admin(self):
        """Optional scripted install: create the first admin from env vars.

        Otherwise there are no accounts until someone completes the first-run
        setup on the login page.
        """
        username = os.environ.get("RIGSHARE_ADMIN_USER")
        password = os.environ.get("RIGSHARE_ADMIN_PASSWORD")
        if self.users or not (username and password):
            if not self.users:
                log.warning("[RigShare] No accounts yet: open ComfyUI in a browser to create the admin account.")
            return
        self.create_user(username, password, perms={"edit": True, "queue": True, "admin": True})
        log.info(f"[RigShare] Created admin '{username}' from environment variables.")

    @property
    def needs_setup(self):
        return not self.users

    def setup_admin(self, username, password, display_name=None):
        """First-run setup: only allowed while no account exists."""
        with self._lock:
            if self.users:
                raise PermissionError("Setup is already complete")
            return self.create_user(username, password, display_name, {"edit": True, "queue": True, "admin": True})

    def _new_user(self, username, password, perms):
        return {
            "username": username,
            "display_name": username,
            "password": _hash_password(password),
            "perms": {p: bool(perms.get(p)) for p in PERMS},
            "created": int(time.time()),
        }

    def _prune_sessions(self):
        now = time.time()
        alive = {t: s for t, s in self.sessions.items() if s["expires"] > now and s["username"] in self.users}
        if len(alive) != len(self.sessions):
            self.sessions = alive
            self._write("sessions.json", self.sessions)

    # ----- config -------------------------------------------------------

    def update_config(self, patch):
        allowed = set(DEFAULT_CONFIG)
        clean = {k: v for k, v in patch.items() if k in allowed}
        self.config = _deep_merge(self.config, clean)
        self._write("config.json", self.config)
        return self.config

    # ----- accounts -----------------------------------------------------

    @staticmethod
    def public_user(user):
        return {
            "username": user["username"],
            "display_name": user.get("display_name") or user["username"],
            "perms": user["perms"],
            "created": user.get("created"),
        }

    def authenticate(self, username, password):
        user = self.users.get(username)
        if user and _verify_password(password, user["password"]):
            return user
        # Burn comparable time so unknown usernames are not distinguishable.
        _verify_password(password, _hash_password("x", b"0" * 16))
        return None

    def _name_taken(self, name, ignore=None):
        """Names differing only by case share one folder on Windows, so they count as the same."""
        return any(u.casefold() == name.casefold() for u in self.users if u != ignore)

    def create_user(self, username, password, display_name=None, perms=None):
        if not USERNAME_RE.match(username or ""):
            raise ValueError("Username must be 2-32 characters: letters, digits, _ . -")
        if self._name_taken(username):
            raise ValueError("Username already exists")
        if len(password or "") < 8:
            raise ValueError("Password must be at least 8 characters")
        user = self._new_user(username, password, perms or {})
        if display_name:
            user["display_name"] = display_name[:32]
        self.users[username] = user
        self._write("users.json", self.users)
        return user

    def rename_user(self, username, new_username):
        """Change a login name; sessions and API keys follow the account."""
        new_username = (new_username or "").strip()
        if new_username == username:
            return self.users[username]
        if username not in self.users:
            raise KeyError(username)
        if not USERNAME_RE.match(new_username):
            raise ValueError("Username must be 2-32 characters: letters, digits, _ . -")
        if self._name_taken(new_username, ignore=username):
            raise ValueError("Username already exists")
        with self._lock:
            user = self.users.pop(username)
            user["username"] = new_username
            if user.get("display_name") in (None, "", username):
                user["display_name"] = new_username
            self.users[new_username] = user
            for session in self.sessions.values():
                if session["username"] == username:
                    session["username"] = new_username
            self._write("users.json", self.users)
            self._write("sessions.json", self.sessions)
        return user

    def update_user(self, username, display_name=None, perms=None, password=None):
        user = self.users.get(username)
        if not user:
            raise KeyError(username)
        if display_name is not None:
            user["display_name"] = display_name.strip()[:32] or username
        if perms is not None:
            new_perms = {p: bool(perms.get(p, user["perms"].get(p))) for p in PERMS}
            if user["perms"].get("admin") and not new_perms["admin"] and self._admin_count() <= 1:
                raise ValueError("Cannot remove the last admin")
            user["perms"] = new_perms
        if password is not None:
            if len(password) < 8:
                raise ValueError("Password must be at least 8 characters")
            user["password"] = _hash_password(password)
            self.revoke_sessions(username)
        self._write("users.json", self.users)
        return user

    def delete_user(self, username):
        user = self.users.get(username)
        if not user:
            raise KeyError(username)
        if user["perms"].get("admin") and self._admin_count() <= 1:
            raise ValueError("Cannot delete the last admin")
        del self.users[username]
        self._write("users.json", self.users)
        self.revoke_sessions(username)

    def _admin_count(self):
        return sum(1 for u in self.users.values() if u["perms"].get("admin"))

    # ----- sessions -----------------------------------------------------

    # Browser logins are JWTs (HS256, 12 h) that reference a server-side
    # session id. The session makes logout, password changes and account
    # deletion take effect immediately instead of waiting for expiry.

    def _issue(self, sid, username):
        now = int(time.time())
        return jwt.encode({"iss": ISSUER, "sub": username, "sid": sid, "iat": now, "nbf": now,
                           "exp": now + TOKEN_TTL}, self.secret)

    def create_session(self, username):
        sid = secrets.token_urlsafe(16)
        self.sessions[sid] = {"username": username, "expires": time.time() + SESSION_TTL, "created": int(time.time())}
        self._write("sessions.json", self.sessions)
        return self._issue(sid, username)

    def _session_claims(self, token, verify_exp=True):
        try:
            claims = jwt.decode(token, self.secret, ISSUER, verify_exp=verify_exp)
        except jwt.InvalidToken:
            return None
        session = self.sessions.get(claims.get("sid", ""))
        if not session or session["expires"] < time.time():
            return None
        # The session (not the token) says who this is, so renames keep sessions valid.
        return {**claims, "sub": session["username"]}

    def user_for_token(self, token):
        """Resolve a session JWT or a personal API key to a user."""
        if not token:
            return None
        if token.startswith(API_KEY_PREFIX):
            return self.user_for_api_key(token)
        claims = self._session_claims(token)
        return self.users.get(claims["sub"]) if claims else None

    def refresh_token(self, token):
        """Return a fresh JWT for the same session when this one is getting old."""
        if not token or token.startswith(API_KEY_PREFIX):
            return None
        claims = self._session_claims(token)
        if not claims or time.time() - claims.get("iat", 0) < TOKEN_REFRESH:
            return None
        return self._issue(claims["sid"], claims["sub"])

    # ----- API keys -----------------------------------------------------

    @staticmethod
    def _key_hash(key):
        return hashlib.sha256(key.encode()).hexdigest()

    def list_api_keys(self, username):
        user = self.users.get(username) or {}
        return [{k: v for k, v in key.items() if k != "hash"} for key in user.get("api_keys", [])]

    def create_api_key(self, username, name):
        user = self.users[username]
        key = API_KEY_PREFIX + secrets.token_urlsafe(32)
        entry = {"id": secrets.token_hex(4), "name": (name or "API key").strip()[:40],
                 "hash": self._key_hash(key), "hint": key[:len(API_KEY_PREFIX) + 4] + "..." + key[-4:],
                 "created": int(time.time()), "last_used": None}
        user.setdefault("api_keys", []).append(entry)
        self._write("users.json", self.users)
        return key, {k: v for k, v in entry.items() if k != "hash"}

    def revoke_api_key(self, username, key_id):
        user = self.users[username]
        before = len(user.get("api_keys", []))
        user["api_keys"] = [k for k in user.get("api_keys", []) if k["id"] != key_id]
        if len(user["api_keys"]) != before:
            self._write("users.json", self.users)
            return True
        return False

    def user_for_api_key(self, key):
        digest = self._key_hash(key)
        for user in self.users.values():
            for entry in user.get("api_keys", []):
                if hmac.compare_digest(entry["hash"], digest):
                    now = int(time.time())
                    if not entry.get("last_used") or now - entry["last_used"] > 300:
                        entry["last_used"] = now
                        self._write("users.json", self.users)
                    return user
        return None

    def revoke_session(self, token):
        claims = self._session_claims(token, verify_exp=False) if token else None
        if claims and self.sessions.pop(claims["sid"], None):
            self._write("sessions.json", self.sessions)

    def revoke_sessions(self, username):
        before = len(self.sessions)
        self.sessions = {t: s for t, s in self.sessions.items() if s["username"] != username}
        if len(self.sessions) != before:
            self._write("sessions.json", self.sessions)

    # ----- guests -------------------------------------------------------

    def guest_perms(self, guest_id):
        base = self.config["guest_permissions"]
        override = self.guests.get(guest_id, {})
        return {"edit": bool(override.get("edit", base.get("edit"))),
                "queue": bool(override.get("queue", base.get("queue"))),
                "admin": False}

    def set_guest_perms(self, guest_id, perms):
        entry = self.guests.setdefault(guest_id, {})
        for p in ("edit", "queue"):
            if p in perms:
                entry[p] = bool(perms[p])
        self._write("guests.json", self.guests)

    # ----- rooms ----------------------------------------------------------

    def _room_path(self, key):
        return os.path.join(self.room_dir, hashlib.sha1(key.encode()).hexdigest()[:20] + ".json")

    def load_rooms(self):
        rooms = []
        for name in os.listdir(self.room_dir):
            if name.endswith(".json"):
                try:
                    with open(os.path.join(self.room_dir, name), "r", encoding="utf-8") as f:
                        rooms.append(json.load(f))
                except (OSError, ValueError):
                    pass
        return rooms

    def save_room(self, room):
        path = self._room_path(room["key"])
        with self._lock:
            with open(path + ".tmp", "w", encoding="utf-8") as f:
                json.dump(room, f)
            os.replace(path + ".tmp", path)

    def delete_room(self, key):
        try:
            os.remove(self._room_path(key))
        except OSError:
            pass

    # ----- snapshots ------------------------------------------------------

    def list_snapshots(self, room=None):
        items = []
        for name in sorted(os.listdir(self.snapshot_dir), reverse=True):
            if not name.endswith(".json"):
                continue
            path = os.path.join(self.snapshot_dir, name)
            try:
                with open(path, "r", encoding="utf-8") as f:
                    meta = json.load(f).get("meta", {})
            except (OSError, ValueError):
                continue
            if room is None or meta.get("room") == room:
                items.append({"id": name[:-5], **meta})
        return items

    # ----- run log (Performance) -----------------------------------------

    RUN_LOG_KEEP = 500

    def runs(self):
        """Finished runs, oldest first: [{id, ms, end, result, workflow, user, name}]."""
        if not hasattr(self, "_runs"):
            self._runs = self._read("runs.json", [])
        return self._runs

    def add_run(self, run):
        with self._lock:
            runs = self.runs()
            runs.append(run)
            del runs[:-self.RUN_LOG_KEEP]
            self._write("runs.json", runs)

    def clear_runs(self, workflow):
        """Forget a workflow's timings (its ETA starts over); returns how many went."""
        with self._lock:
            runs = self.runs()
            kept = [r for r in runs if r.get("workflow") != workflow]
            gone = len(runs) - len(kept)
            self._runs = kept
            self._write("runs.json", kept)
            return gone

    def add_snapshot(self, doc, author, room, room_name, label=None, kind=None):
        """``kind`` "save": taken when the file was saved; those rotate like automatic ones."""
        # Sorted by name, newest first: include microseconds so saves in the same second keep their order.
        now = time.time()
        snap_id = time.strftime("%Y%m%d-%H%M%S", time.localtime(now)) + f"-{int(now % 1 * 1e6):06d}-" + secrets.token_hex(2)
        meta = {
            "time": int(time.time()),
            "author": author,
            "label": (label or "")[:80],
            "room": room,
            "room_name": room_name,
            "nodes": len(doc.get("nodes", [])) if isinstance(doc, dict) else 0,
        }
        if kind:
            meta["kind"] = kind
        with open(os.path.join(self.snapshot_dir, snap_id + ".json"), "w", encoding="utf-8") as f:
            json.dump({"meta": meta, "doc": doc}, f)
        self._prune_snapshots(room)
        return {"id": snap_id, **meta}

    def get_snapshot(self, snap_id):
        if not re.match(r"^[\w-]+$", snap_id or ""):
            return None
        try:
            with open(os.path.join(self.snapshot_dir, snap_id + ".json"), "r", encoding="utf-8") as f:
                return json.load(f)
        except (OSError, ValueError):
            return None

    def _prune_snapshots(self, room):
        keep = max(1, int(self.config.get("snapshot_keep", 30)))
        # Named (manual) snapshots are kept; automatic and on-save ones rotate, each on its own.
        snaps = self.list_snapshots(room)
        auto = [s for s in snaps if not s.get("label") and not s.get("kind")]
        saves = [s for s in snaps if s.get("kind") == "save"]
        for snap in auto[keep:] + saves[keep:]:
            try:
                os.remove(os.path.join(self.snapshot_dir, snap["id"] + ".json"))
            except OSError:
                pass