#!/usr/bin/env python3
"""
overleaf-ce-sync — two-way sync CLI for a self-hosted Overleaf Community Edition.

It logs in over plain HTTP (CE has no CAPTCHA), pulls a project as a zip and
pushes individual files back via the upload endpoint. The only thing that needs
socket.io is discovering the project's root folder id, which is required by the
upload endpoint. It uses xhr-polling (plain HTTP, no websocket upgrade) and the
result is cached, so socket.io is hit at most once per project.

Endpoints used (verified against the CE source):
  GET  /login                         -> session cookie + CSRF token
  POST /login                         -> authenticate (no CAPTCHA on CE)
  GET  /dev/csrf                      -> CSRF token for the current session
  GET  /user/projects                 -> { projects: [{_id, name, accessLevel}] }
  GET  /Project/<id>/download/zip     -> project zip (pull)
  POST /Project/<id>/upload?folder_id=<root>  (field: qqfile)  -> push a file
  GET  /socket.io/1/xhr-polling/...  -> joinProjectResponse -> rootFolder[0]._id

Dependencies: requests
"""

import argparse
import getpass
import hashlib
import json
import os
import subprocess
import sys
import tempfile
import time
import zipfile
from fnmatch import fnmatch
from io import BytesIO
from urllib.parse import urlparse

try:
    import requests
except ImportError:
    sys.exit("Missing dependency: pip install requests")

CONFIG_DIR = os.path.expanduser("~/.config/overleaf-ce")
# CE's default session cookie is 'overleaf.sid' (settings.defaults.js: COOKIE_NAME).
# overleaf.com uses 'overleaf_session2'. We don't hardcode either: login success is
# detected from the response, and we persist whatever cookies the server sets.
SESSION_COOKIE = os.environ.get("OVERLEAF_CE_COOKIE_NAME", "overleaf.sid")
PROJECT_FILE = ".overleaf-ce.json"
MANIFEST_FILE = ".overleaf-ce-manifest.json"

# Always-ignored, in addition to whatever .gitignore / .olignore say.
DEFAULT_IGNORES = [
    ".git/", ".git", ".gitignore", PROJECT_FILE, MANIFEST_FILE, ".olignore",
    "*.pyc", "__pycache__/", ".DS_Store",
]


# --------------------------------------------------------------------------- #
# small helpers
# --------------------------------------------------------------------------- #
def die(msg, code=1):
    print(f"error: {msg}", file=sys.stderr)
    sys.exit(code)


def info(msg):
    print(msg)


def normalize_host(host):
    """Accept 'overleaf.all.ar' or 'https://overleaf.all.ar' -> base URL."""
    host = host.strip().rstrip("/")
    if not host.startswith(("http://", "https://")):
        host = "https://" + host
    return host


def cookie_path(base):
    safe = urlparse(base).netloc.replace(":", "_")
    return os.path.join(CONFIG_DIR, f"cookies_{safe}.json")


def save_cookies(session, base):
    os.makedirs(CONFIG_DIR, exist_ok=True)
    path = cookie_path(base)
    with open(path, "w") as fh:
        json.dump(requests.utils.dict_from_cookiejar(session.cookies), fh)
    os.chmod(path, 0o600)


def load_cookies(session, base):
    path = cookie_path(base)
    if not os.path.exists(path):
        return False
    with open(path) as fh:
        session.cookies.update(requests.utils.cookiejar_from_dict(json.load(fh)))
    return len(session.cookies) > 0


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 16), b""):
            h.update(chunk)
    return h.hexdigest()


# --------------------------------------------------------------------------- #
# git helpers — push syncs *committed* state, so it talks to git, not the disk
# --------------------------------------------------------------------------- #
def _run_git(directory, *args, want_bytes=False):
    return subprocess.run(
        ["git", *args], cwd=directory,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=not want_bytes,
    )


def is_git_repo(directory):
    return _run_git(directory, "rev-parse", "--git-dir").returncode == 0


def git_head(directory):
    r = _run_git(directory, "rev-parse", "HEAD")
    return r.stdout.strip() if r.returncode == 0 else None


def git_is_dirty(directory, include_untracked=True):
    args = ["status", "--porcelain"]
    if not include_untracked:
        args.append("--untracked-files=no")
    return bool(_run_git(directory, *args).stdout.strip())


def git_tracked_files(directory, commit="HEAD"):
    r = _run_git(directory, "ls-tree", "-r", "--name-only", "-z", commit)
    return [p for p in r.stdout.split("\0") if p]


def git_diff_files(directory, base, head, diff_filter):
    r = _run_git(directory, "diff", "--name-only", "-z",
                 f"--diff-filter={diff_filter}", base, head)
    if r.returncode != 0:
        die(f"git diff failed: {r.stderr.strip()}")
    return [p for p in r.stdout.split("\0") if p]


def git_show_bytes(directory, commit, path):
    """Bytes of a file as committed at <commit> (not the working-tree copy)."""
    r = _run_git(directory, "show", f"{commit}:{path}", want_bytes=True)
    if r.returncode != 0:
        die(f"git show {commit}:{path} failed: "
            f"{r.stderr.decode(errors='replace').strip()}")
    return r.stdout


def ensure_gitignore(directory):
    """Keep our local sync metadata out of the user's commits."""
    path = os.path.join(directory, ".gitignore")
    needed = [PROJECT_FILE, MANIFEST_FILE]
    existing = ""
    if os.path.exists(path):
        with open(path) as fh:
            existing = fh.read()
    missing = [n for n in needed if n not in existing.split()]
    if missing:
        with open(path, "a") as fh:
            if existing and not existing.endswith("\n"):
                fh.write("\n")
            fh.write("\n".join(missing) + "\n")


def commit_sync(directory, message):
    """Stage everything (minus our metadata) and commit. Returns HEAD sha."""
    ensure_gitignore(directory)
    _run_git(directory, "add", "-A")
    _run_git(directory, "commit", "-q", "-m", message)  # no-op if nothing staged
    return git_head(directory)


# --------------------------------------------------------------------------- #
# project config (per working directory)
# --------------------------------------------------------------------------- #
def read_project_cfg(directory):
    path = os.path.join(directory, PROJECT_FILE)
    if os.path.exists(path):
        with open(path) as fh:
            return json.load(fh)
    return {}


def write_project_cfg(directory, cfg):
    with open(os.path.join(directory, PROJECT_FILE), "w") as fh:
        json.dump(cfg, fh, indent=2)


def read_manifest(directory):
    path = os.path.join(directory, MANIFEST_FILE)
    if os.path.exists(path):
        with open(path) as fh:
            return json.load(fh)
    return {}


def write_manifest(directory, manifest):
    with open(os.path.join(directory, MANIFEST_FILE), "w") as fh:
        json.dump(manifest, fh, indent=2)


# --------------------------------------------------------------------------- #
# auth / csrf
# --------------------------------------------------------------------------- #
def get_csrf(session, base):
    # /dev/csrf returns the token in plain text for the current session.
    try:
        r = session.get(f"{base}/dev/csrf", timeout=20)
        if r.ok and "<" not in r.text and len(r.text) < 200:
            return r.text.strip()
    except requests.RequestException:
        pass
    # Fallback: scrape it from a rendered page.
    r = session.get(f"{base}/login", timeout=20)
    for needle in ('window.csrfToken = "', 'name="ol-csrfToken" content="'):
        i = r.text.find(needle)
        if i != -1:
            start = i + len(needle)
            return r.text[start:r.text.index('"', start)]
    die("could not obtain CSRF token (is this an Overleaf instance?)")


def make_session(base, require_auth=True):
    session = requests.Session()
    session.headers["User-Agent"] = "overleaf-ce-sync/1.0"
    # CE decides JSON-vs-redirect from Accept; force JSON so success/expiry are
    # unambiguous (otherwise login redirects and failures look like successes).
    session.headers["Accept"] = "application/json"
    have = load_cookies(session, base)
    if require_auth and not have:
        die("not logged in for this host — run: overleaf-ce-sync login")
    return session


def cmd_login(args):
    base = normalize_host(args.host or input("Overleaf host (e.g. overleaf.all.ar): "))
    email = args.email or input("email: ")
    password = args.password or getpass.getpass("password: ")

    session = requests.Session()
    session.headers["User-Agent"] = "overleaf-ce-sync/1.0"
    session.headers["Accept"] = "application/json"  # get {redir} JSON, not a 302
    csrf = get_csrf(session, base)
    r = session.post(
        f"{base}/login",
        json={"_csrf": csrf, "email": email, "password": password},
        headers={"X-CSRF-Token": csrf},
        timeout=30,
        allow_redirects=False,
    )
    # CE returns 200 {redir} on success, 401 {message} on bad credentials.
    if r.status_code == 200:
        if not session.cookies:
            die("login returned 200 but the server set no cookie — unexpected")
        save_cookies(session, base)
        # Remember the last host as the default for bare commands.
        os.makedirs(CONFIG_DIR, exist_ok=True)
        with open(os.path.join(CONFIG_DIR, "config.json"), "w") as fh:
            json.dump({"default_host": base}, fh)
        info(f"logged in to {base} as {email} "
             f"(cookie: {', '.join(c.name for c in session.cookies)})")
    else:
        msg = ""
        try:
            msg = r.json().get("message", {})
            msg = msg.get("text", msg) if isinstance(msg, dict) else msg
        except ValueError:
            pass
        die(f"login failed ({r.status_code}) {msg}")


def default_host(explicit):
    if explicit:
        return normalize_host(explicit)
    path = os.path.join(CONFIG_DIR, "config.json")
    if os.path.exists(path):
        with open(path) as fh:
            return json.load(fh).get("default_host")
    return None


def resolve_base(args, cfg):
    """Host precedence: explicit --host > per-project config > global default."""
    if args.host:
        return normalize_host(args.host)
    return (cfg.get("host") or default_host(None)
            or die("no host — run: overleaf-ce-sync login"))


# --------------------------------------------------------------------------- #
# list
# --------------------------------------------------------------------------- #
def cmd_list(args):
    base = default_host(args.host) or die("no host — run: overleaf-ce-sync login")
    session = make_session(base)
    r = session.get(f"{base}/user/projects", timeout=30)
    if r.status_code == 302 or r.url.endswith("/login"):
        die("session expired — run: overleaf-ce-sync login")
    r.raise_for_status()
    projects = r.json().get("projects", [])
    if not projects:
        info("(no projects)")
        return
    width = max(len(p["_id"]) for p in projects)
    for p in projects:
        info(f"{p['_id']:<{width}}  {p['name']}")


# --------------------------------------------------------------------------- #
# pull
# --------------------------------------------------------------------------- #
def cmd_pull(args):
    directory = os.path.abspath(args.dir)
    cfg = read_project_cfg(directory)
    base = resolve_base(args, cfg)
    project_id = args.project_id or cfg.get("project_id") or die("no project id")

    if is_git_repo(directory) and git_is_dirty(directory) and not args.force:
        die("working tree has uncommitted changes — commit them first, or pass "
            "--force to overwrite local files with the remote copy")

    session = make_session(base)
    info(f"pulling {project_id} from {base} ...")
    r = session.get(f"{base}/Project/{project_id}/download/zip", timeout=120)
    if r.url.endswith("/login"):
        die("session expired — run: overleaf-ce-sync login")
    r.raise_for_status()
    if "zip" not in r.headers.get("Content-Type", "") and not r.content[:2] == b"PK":
        die("did not receive a zip (wrong project id or no access?)")

    os.makedirs(directory, exist_ok=True)
    zf = zipfile.ZipFile(BytesIO(r.content))
    names = [n for n in zf.namelist() if not n.endswith("/")]
    if args.prune:
        _prune_local(directory, names)
    zf.extractall(directory)
    info(f"extracted {len(names)} file(s) into {directory}")

    cfg.update({"host": base, "project_id": project_id})

    # In a git repo, snapshot the pulled state as a sync commit. push then diffs
    # against this commit, so it only ever sends *committed* changes you make
    # afterwards (and re-pulling without committing never re-pushes the remote).
    if is_git_repo(directory):
        sync_commit = commit_sync(directory, f"overleaf pull {project_id}")
        if sync_commit:
            cfg["last_sync_commit"] = sync_commit
            info(f"synced at commit {sync_commit[:8]}")
    write_project_cfg(directory, cfg)


def _prune_local(directory, keep_names):
    keep = set(keep_names)
    for root, _dirs, files in os.walk(directory):
        for f in files:
            full = os.path.join(root, f)
            rel = os.path.relpath(full, directory).replace(os.sep, "/")
            if rel in keep or is_ignored(rel, directory):
                continue
            os.remove(full)


# --------------------------------------------------------------------------- #
# ignore handling
# --------------------------------------------------------------------------- #
def load_ignore_patterns(directory):
    patterns = list(DEFAULT_IGNORES)
    for name in (".gitignore", ".olignore"):
        path = os.path.join(directory, name)
        if os.path.exists(path):
            with open(path) as fh:
                for line in fh:
                    line = line.strip()
                    if line and not line.startswith("#"):
                        patterns.append(line)
    return patterns


def is_ignored(rel, directory, patterns=None):
    if patterns is None:
        patterns = load_ignore_patterns(directory)
    base = os.path.basename(rel)
    for pat in patterns:
        p = pat.rstrip("/")
        if fnmatch(rel, p) or fnmatch(base, p):
            return True
        # directory prefix match, e.g. "build/" ignores build/**
        if rel == p or rel.startswith(p + "/"):
            return True
    return False


def iter_local_files(directory):
    patterns = load_ignore_patterns(directory)
    for root, dirs, files in os.walk(directory):
        # prune ignored directories in-place for speed
        rel_root = os.path.relpath(root, directory).replace(os.sep, "/")
        dirs[:] = [
            d for d in dirs
            if not is_ignored(
                (f"{rel_root}/{d}" if rel_root != "." else d), directory, patterns
            )
        ]
        for f in files:
            full = os.path.join(root, f)
            rel = os.path.relpath(full, directory).replace(os.sep, "/")
            if not is_ignored(rel, directory, patterns):
                yield rel, full


# --------------------------------------------------------------------------- #
# discover the root folder id via socket.io 0.9 xhr-polling (plain HTTP).
#
# We use xhr-polling rather than a real websocket on purpose: it is pure HTTP and
# traverses any reverse proxy. Forcing a websocket upgrade 502s behind proxies
# that don't pass Upgrade/Connection headers (Overleaf itself falls back to
# polling in that case). On connect the server pushes a `joinProjectResponse`
# event containing the project tree, from which we read rootFolder[0]._id.
# --------------------------------------------------------------------------- #
def _ms():
    return int(time.time() * 1000)


def _decode_payload(text):
    """Split a socket.io 0.9 xhr-polling body into individual frames."""
    if not text:
        return []
    if text[0] != "�":          # single, unframed message
        return [text]
    frames, i, n = [], 0, len(text)
    while i < n:
        j = text.index("�", i + 1)
        length = int(text[i + 1:j])
        frames.append(text[j + 1:j + 1 + length])
        i = j + 1 + length
    return frames


def discover_root_folder_id(base, session, project_id):
    params = {"projectId": project_id, "t": _ms()}
    # 1) handshake -> session id
    h = session.get(f"{base}/socket.io/1/", params=params, timeout=20)
    h.raise_for_status()
    sid = h.text.split(":", 1)[0]
    poll = f"{base}/socket.io/1/xhr-polling/{sid}"

    deadline = time.time() + 30
    while time.time() < deadline:
        r = session.get(poll, params={"projectId": project_id, "t": _ms()}, timeout=35)
        if r.status_code != 200:
            die(f"socket.io polling failed: HTTP {r.status_code} "
                "(session expired? try: overleaf-ce-sync login)")
        for frame in _decode_payload(r.text):
            if frame.startswith("2:"):                     # heartbeat -> ack
                session.post(poll, params={"projectId": project_id, "t": _ms()},
                             data="2::", timeout=20)
            elif frame.startswith("5:"):                   # event: [5,id,ep,json]
                payload = json.loads(frame.split(":", 3)[3])
                if payload.get("name") == "joinProjectResponse":
                    return payload["args"][0]["project"]["rootFolder"][0]["_id"]
                if payload.get("name") == "connectionRejected":
                    die(f"connection rejected: {payload.get('args')}")
            elif frame.startswith("0:"):                   # server disconnect
                die("socket.io disconnected before joinProjectResponse "
                    "(session expired or no access to this project?)")
    die("timed out waiting for joinProjectResponse over xhr-polling\n"
        "       fallback: pass --root-folder-id <id> (see README)")


def resolve_root_folder_id(base, session, project_id, directory, override=None):
    cfg = read_project_cfg(directory)
    if override:
        rid = override
    elif cfg.get("root_folder_id"):
        return cfg["root_folder_id"]
    else:
        info("discovering root folder id via socket.io xhr-polling ...")
        rid = discover_root_folder_id(base, session, project_id)
    cfg.update({"host": base, "project_id": project_id, "root_folder_id": rid})
    write_project_cfg(directory, cfg)
    return rid


# --------------------------------------------------------------------------- #
# push
# --------------------------------------------------------------------------- #
def cmd_push(args):
    directory = os.path.abspath(args.dir)
    cfg = read_project_cfg(directory)
    base = resolve_base(args, cfg)
    project_id = args.project_id or cfg.get("project_id") or die("no project id")

    # push publishes the *committed* state (HEAD), never the working tree.
    if not is_git_repo(directory):
        die("push needs a git repo (it pushes committed state). use "
            "'overleaf-ce-sync clone', or run: git init && git add -A && git commit")
    head = git_head(directory)
    if not head:
        die("no commits yet — commit your files first: git add -A && git commit")
    if git_is_dirty(directory, include_untracked=False):
        info("note: tracked files have uncommitted changes; pushing committed state (HEAD) only")

    last_sync = cfg.get("last_sync_commit")
    if last_sync == head:
        info("nothing to push (HEAD already synced)")
        return

    if last_sync:
        changed = git_diff_files(directory, last_sync, head, "ACMR")
        deleted = git_diff_files(directory, last_sync, head, "D")
    else:
        changed = git_tracked_files(directory, head)   # first push: send everything
        deleted = []

    # never push our own metadata files
    skip = {PROJECT_FILE, MANIFEST_FILE, ".gitignore", ".olignore"}
    changed = [p for p in changed if p not in skip]
    deleted = [p for p in deleted if p not in skip]

    if deleted:
        info(f"note: {len(deleted)} file(s) deleted in your commits are NOT removed remotely:")
        for d in deleted:
            info(f"  - {d}")

    if not changed:
        info("nothing to push (no committed file changes since last sync)")
        cfg["last_sync_commit"] = head
        write_project_cfg(directory, cfg)
        return

    session = make_session(base)
    csrf = get_csrf(session, base)
    root_id = resolve_root_folder_id(
        base, session, project_id, directory, override=args.root_folder_id
    )
    cfg = read_project_cfg(directory)  # re-read: resolve may have cached root_folder_id

    for rel in changed:
        content = git_show_bytes(directory, head, rel)
        with tempfile.NamedTemporaryFile() as tmp:
            tmp.write(content)
            tmp.flush()
            _upload_file(session, base, project_id, root_id, csrf, rel, tmp.name)
        info(f"  pushed {rel}")

    cfg["last_sync_commit"] = head
    write_project_cfg(directory, cfg)
    info(f"pushed {len(changed)} file(s) to {project_id} (commit {head[:8]})")


def _upload_file(session, base, project_id, root_id, csrf, rel, full):
    name = os.path.basename(rel)
    # nested file -> send its path so the server mkdirs; root file -> "null"
    relative_path = rel if "/" in rel else "null"
    with open(full, "rb") as fh:
        r = session.post(
            f"{base}/Project/{project_id}/upload",
            params={"folder_id": root_id},
            data={"name": name, "relativePath": relative_path, "_csrf": csrf},
            files={"qqfile": (name, fh)},
            headers={"X-CSRF-Token": csrf},
            timeout=120,
        )
    if not r.ok:
        die(f"upload failed for {rel}: HTTP {r.status_code}")
    try:
        body = r.json()
    except ValueError:
        die(f"upload failed for {rel}: unexpected response")
    if not body.get("success"):
        die(f"upload failed for {rel}: {body.get('error', body)}")


# --------------------------------------------------------------------------- #
# clone = pull + git init/commit
# --------------------------------------------------------------------------- #
def cmd_clone(args):
    directory = os.path.abspath(args.dir or args.project_id)
    os.makedirs(directory, exist_ok=True)
    pull_args = argparse.Namespace(
        host=args.host, project_id=args.project_id, dir=directory,
        prune=False, force=False,
    )
    cmd_pull(pull_args)
    if not is_git_repo(directory):
        _run_git(directory, "init", "-q")
    sync_commit = commit_sync(directory, f"overleaf pull {args.project_id}")
    cfg = read_project_cfg(directory)
    if sync_commit:
        cfg["last_sync_commit"] = sync_commit
        write_project_cfg(directory, cfg)
    info(f"cloned into {directory} "
         f"(git repo, synced at {sync_commit[:8] if sync_commit else '?'})")


# --------------------------------------------------------------------------- #
# argparse
# --------------------------------------------------------------------------- #
def build_parser():
    p = argparse.ArgumentParser(prog="overleaf-ce-sync", description=__doc__.splitlines()[1])
    sub = p.add_subparsers(dest="cmd", required=True)

    sp = sub.add_parser("login", help="authenticate and store the session cookie")
    sp.add_argument("--host"); sp.add_argument("--email"); sp.add_argument("--password")
    sp.set_defaults(func=cmd_login)

    sp = sub.add_parser("list", help="list your projects (id + name)")
    sp.add_argument("--host")
    sp.set_defaults(func=cmd_list)

    sp = sub.add_parser("pull", help="download a project, overwrite local files, commit a sync point")
    sp.add_argument("project_id", nargs="?")
    sp.add_argument("--host"); sp.add_argument("-d", "--dir", default=".")
    sp.add_argument("--prune", action="store_true",
                    help="delete local files that no longer exist on Overleaf")
    sp.add_argument("--force", action="store_true",
                    help="overwrite even if the working tree has uncommitted changes")
    sp.set_defaults(func=cmd_pull)

    sp = sub.add_parser("push", help="upload committed changes (HEAD) to a project")
    sp.add_argument("project_id", nargs="?")
    sp.add_argument("--host"); sp.add_argument("-d", "--dir", default=".")
    sp.add_argument("--root-folder-id",
                    help="skip socket.io discovery (get it once from mongo, see README)")
    sp.set_defaults(func=cmd_push)

    sp = sub.add_parser("clone", help="pull into a new dir and init a git repo")
    sp.add_argument("project_id")
    sp.add_argument("dir", nargs="?")
    sp.add_argument("--host")
    sp.set_defaults(func=cmd_clone)

    return p


def main():
    args = build_parser().parse_args()
    try:
        args.func(args)
    except requests.RequestException as e:
        die(f"network error: {e}")
    except KeyboardInterrupt:
        sys.exit(130)


if __name__ == "__main__":
    main()