#!/usr/bin/env python3
"""Kokkoja local program.

Serves the Kokkoja client to a web browser on this computer, and reads
and writes the user's notebooks for it.  A notebook is a directory; a page
is one file in it.  It also lets the client's own file chooser look through
the user's directories for image files, and writes an exported picture of
a page where the user chooses, neither of which a web page can do itself.

The client must be loaded from this computer (or over HTTPS) because
browsers allow their built-in cryptography only there.

Only the user's own browser may use the notebooks.  The program listens on
this computer alone, and every notebook request must carry a secret token.
The browser gets the token from the address this program opens or prints;
another user of the computer, or a web page from elsewhere, does not have it.

Uses only the Python standard library.
"""

import argparse
import functools
import hmac
import http.server
import json
import os
import pathlib
import re
import secrets
import subprocess
import sys
import threading
import time
import urllib.parse
import webbrowser

CLIENT_DIRECTORY = pathlib.Path(__file__).resolve().parent / "client"

MAX_PAGE_BYTES = 64 << 20

# The page being written in a notebook, and the pages kept from it.
CURRENT_PAGE = "Current.svg"
KEPT_PAGE = re.compile(r"^(\d{4}_\d{2}_\d{2})_(\d{2,})\.svg$")

# Notebook files are changed by one request at a time.
notebook_lock = threading.Lock()

# A notebook or page name: one path component, not hidden, no control
# characters, no characters that a file system might refuse.
NAME = re.compile(r'^[^\x00-\x1f/\\:*?"<>|.][^\x00-\x1f/\\:*?"<>|]{0,119}$')

# The image files the file chooser offers, and what a browser is told they
# are.  The chooser reads these and nothing else.
IMAGE_TYPES = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
               ".gif": "image/gif", ".webp": "image/webp", ".bmp": "image/bmp"}
MAX_IMAGE_FILE_BYTES = 64 << 20

# An exported page is a PNG file, and nothing else is written outside the
# notebooks: the name must end in .png and the bytes must begin like this.
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"

LOCAL_HOST = re.compile(r"^(localhost|127\.0\.0\.1|\[::1\])(:\d+)?$")
LOCAL_ORIGIN = re.compile(r"^http://(localhost|127\.0\.0\.1|\[::1\])(:\d+)?$")


def config_directory():
    """Where Kokkoja keeps its own settings, by each system's convention."""
    if sys.platform == "win32":
        base = os.environ.get("APPDATA") or pathlib.Path.home() / "AppData" / "Roaming"
    elif sys.platform == "darwin":
        base = pathlib.Path.home() / "Library" / "Application Support"
    else:
        base = os.environ.get("XDG_CONFIG_HOME") or pathlib.Path.home() / ".config"
    return pathlib.Path(base) / "kokkoja"


def default_notebooks_directory():
    """A Kokkoja directory in the user's documents directory."""
    documents = pathlib.Path.home() / "Documents"
    if sys.platform.startswith("linux"):
        # The documents directory has another name in other languages.
        try:
            found = subprocess.run(["xdg-user-dir", "DOCUMENTS"], capture_output=True,
                                   text=True, timeout=5).stdout.strip()
            if found and pathlib.Path(found) != pathlib.Path.home():
                documents = pathlib.Path(found)
        except (OSError, subprocess.SubprocessError):
            pass
    return documents / "Kokkoja"


def load_token(directory):
    """The secret token, made on first use and kept in a private file."""
    path = directory / "token"
    try:
        token = path.read_text().strip()
        if token:
            return token
    except OSError:
        pass
    token = secrets.token_urlsafe(24)
    path.parent.mkdir(parents=True, exist_ok=True)
    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(descriptor, "w") as file:
        file.write(token + "\n")
    return token


def write_atomically(path, data):
    """Replace a file so that it is never seen half written."""
    temporary = path.with_name("." + path.name + ".writing")
    with open(temporary, "wb") as file:
        file.write(data)
        file.flush()
        os.fsync(file.fileno())
    os.replace(temporary, path)


def keep_page(notebook):
    """Keep a copy of the notebook's current page as a dated page.

    The copy is named for the day the page was last changed, with a serial
    number: 2026_10_02_01.svg.  Nothing is copied if the page is blank or
    if an identical page is already kept.

    Returns (name, new): the name of the kept page, or None if there was
    nothing to keep, and whether the copy was made just now.
    """
    current = notebook / CURRENT_PAGE
    try:
        data = current.read_bytes()
        changed = current.stat().st_mtime
    except FileNotFoundError:
        return None, False
    if b"data-stroke=" not in data and b"<image" not in data:
        return None, False

    kept = [entry for entry in notebook.iterdir()
            if KEPT_PAGE.match(entry.name) and entry.is_file()]
    for entry in kept:
        if entry.stat().st_size == len(data) and entry.read_bytes() == data:
            return entry.name, False

    date = time.strftime("%Y_%m_%d", time.localtime(changed))
    serials = [int(match.group(2)) for match in map(KEPT_PAGE.match, (e.name for e in kept))
               if match.group(1) == date]
    name = f"{date}_{max(serials, default=0) + 1:02d}.svg"
    write_atomically(notebook / name, data)
    os.utime(notebook / name, (changed, changed))
    return name, True


MAX_IDENTITY_BYTES = 4096


def site_of(relay, site):
    """The Kokkoja web site that invitations link to, or "".

    Given explicitly, it is used as given.  Otherwise it is taken from the
    relay's address when the relay is a Kokkoja server's, at wss://<host>/relay:
    such a server serves the client at https://<host>/ as well.
    """
    if site:
        return site.rstrip("/") + "/"
    parts = urllib.parse.urlsplit(relay)
    if parts.scheme == "wss" and parts.path.rstrip("/") == "/relay":
        return f"https://{parts.netloc}/"
    return ""


class Handler(http.server.SimpleHTTPRequestHandler):
    relay = ""
    site = ""
    token = ""
    notebooks = pathlib.Path()
    config = pathlib.Path()

    # Replies.

    def send_json(self, value, status=200):
        self.send_body(json.dumps(value).encode(), "application/json", status)

    def send_body(self, body, content_type, status=200):
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def refuse(self, status, reason):
        self.send_json({"error": reason}, status)

    def end_headers(self):
        # The client files change during development; never serve stale ones.
        self.send_header("Cache-Control", "no-store")
        super().end_headers()

    def log_message(self, format, *arguments):
        pass

    # Requests.

    def do_GET(self):
        path = self.path.split("?")[0]
        if path == "/config.json":
            # How the client was served: by this program, which keeps the
            # notebooks as files.  A web site serves its own config.json
            # with "storage": "browser".
            self.send_json({"relay": self.relay, "storage": "files", "site": self.site})
        elif path == "/api/identity":
            self.identity_request("GET")
        elif path == "/api/files":
            self.files_request()
        elif path == "/api/file":
            self.file_request("GET")
        elif path == "/api" or path.startswith("/api/"):
            self.notebook_request("GET", path)
        else:
            super().do_GET()

    def do_PUT(self):
        path = self.path.split("?")[0]
        if path == "/api/identity":
            self.identity_request("PUT")
        elif path == "/api/file":
            self.file_request("PUT")
        else:
            self.notebook_request("PUT", path)

    def do_POST(self):
        self.notebook_request("POST", self.path.split("?")[0])

    def allowed(self):
        """Whether a notebook request comes from the user's own browser."""
        if not LOCAL_HOST.match(self.headers.get("Host", "")):
            self.refuse(403, "not addressed to this computer")
            return False
        origin = self.headers.get("Origin")
        if origin is not None and not LOCAL_ORIGIN.match(origin):
            self.refuse(403, "request from another web site")
            return False
        given = self.headers.get("Authorization", "")
        if not hmac.compare_digest(given.encode(), ("Bearer " + self.token).encode()):
            self.refuse(401, "missing or wrong token")
            return False
        return True

    def identity_request(self, method):
        """
        GET /api/identity   the user's identity (signing keys), or 404
        PUT /api/identity   store it, once; a second PUT is refused

        The identity is kept here rather than in the browser so that the
        user is the same participant in every session, from any browser on
        this computer, and their writing from an earlier lesson stays
        theirs.  The file is private to the user.
        """
        if not self.allowed():
            return
        path = self.config / "identity.json"
        try:
            if method == "GET":
                self.send_body(path.read_bytes(), "application/json")
                return
            length = int(self.headers.get("Content-Length", "-1"))
            if not 0 < length <= MAX_IDENTITY_BYTES:
                return self.refuse(413, "not an identity")
            data = self.rfile.read(length)
            try:
                json.loads(data)
            except ValueError:
                return self.refuse(400, "not an identity")
            with notebook_lock:
                if path.exists():
                    return self.refuse(409, "an identity is already kept")
                path.parent.mkdir(parents=True, exist_ok=True)
                descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
                with os.fdopen(descriptor, "wb") as file:
                    file.write(data)
            self.send_json({"kept": True})
        except FileNotFoundError:
            self.refuse(404, "no identity yet")
        except (OSError, ValueError) as error:
            self.refuse(500, str(error))

    def notebook_request(self, method, path):
        """
        GET /api/notebooks                     the notebooks
        GET /api/notebooks/<notebook>          the pages of a notebook
        GET /api/notebooks/<notebook>/<page>   a page
        PUT /api/notebooks/<notebook>/<page>   write a page
        POST /api/notebooks/<notebook>/keep    keep a copy of Current.svg
        """
        if not self.allowed():
            return
        parts = [urllib.parse.unquote(part) for part in path.split("/")[2:]]
        if not parts or parts[0] != "notebooks" or len(parts) > 3:
            return self.refuse(404, "no such request")
        names = parts[1:]
        if not all(NAME.match(name) for name in names):
            return self.refuse(400, "name not allowed")
        keeping = method == "POST" and len(names) == 2 and names[1] == "keep"
        if len(names) == 2 and not keeping and not names[1].endswith(".svg"):
            return self.refuse(400, "a page name ends in .svg")
        target = self.notebooks.joinpath(*names)

        try:
            if keeping:
                with notebook_lock:
                    name, new = keep_page(target.parent) if target.parent.is_dir() else (None, False)
                self.send_json({"kept": name, "new": new})
            elif method == "POST":
                self.refuse(405, "nothing to do")
            elif method == "PUT":
                if len(names) != 2:
                    return self.refuse(405, "only a page can be written")
                length = int(self.headers.get("Content-Length", "-1"))
                if not 0 <= length <= MAX_PAGE_BYTES:
                    return self.refuse(413, "page too large")
                data = self.rfile.read(length)
                with notebook_lock:
                    target.parent.mkdir(parents=True, exist_ok=True)
                    write_atomically(target, data)
                self.send_json({"saved": names[1]})
            elif len(names) == 0:
                self.send_json({"directory": str(self.notebooks),
                                "notebooks": self.listing(self.notebooks, directories=True)})
            elif len(names) == 1:
                self.send_json({"pages": self.listing(target, directories=False)})
            else:
                self.send_body(target.read_bytes(), "image/svg+xml")
        except FileNotFoundError:
            self.refuse(404, "not found")
        except (OSError, ValueError) as error:
            self.refuse(500, str(error))

    def query(self):
        """The named values after '?' in the request's address."""
        values = urllib.parse.parse_qs(urllib.parse.urlsplit(self.path).query,
                                       keep_blank_values=True)
        return {name: value[-1] for name, value in values.items()}

    def files_request(self):
        """
        GET /api/files?directory=<full path>   the directories and the image
                                               files in a directory

        Without a directory: the notebooks directory, or the user's home
        directory if there are no notebooks yet.  Hidden entries are left
        out.  Each entry comes with its full path, so that the client never
        has to put a path together.
        """
        if not self.allowed():
            return
        given = self.query().get("directory", "")
        try:
            if given:
                directory = pathlib.Path(given).expanduser()
                if not directory.is_absolute():
                    return self.refuse(400, "not a full path")
            else:
                directory = self.notebooks if self.notebooks.is_dir() else pathlib.Path.home()
            directory = pathlib.Path(os.path.abspath(directory))
            directories = []
            files = []
            for entry in directory.iterdir():
                if entry.name.startswith("."):
                    continue
                try:
                    if entry.is_dir():
                        directories.append({"name": entry.name, "path": str(entry)})
                    elif entry.is_file() and entry.suffix.lower() in IMAGE_TYPES:
                        files.append({"name": entry.name, "path": str(entry),
                                      "size": entry.stat().st_size})
                except OSError:
                    continue
            by_name = lambda entry: (entry["name"].casefold(), entry["name"])
            self.send_json({
                "directory": str(directory),
                "parent": None if directory.parent == directory else str(directory.parent),
                "directories": sorted(directories, key=by_name),
                "files": sorted(files, key=by_name)})
        except FileNotFoundError:
            self.refuse(404, "no such directory")
        except NotADirectoryError:
            self.refuse(400, "not a directory")
        except PermissionError:
            self.refuse(403, "this directory may not be read")
        except (OSError, ValueError) as error:
            self.refuse(500, str(error))

    def file_request(self, method):
        """
        GET /api/file?path=<full path>               an image file
        PUT /api/file?directory=<full path>&name=<name>.png[&replace=1]
                                                     write an exported page

        Only image files are read, by their names.  Only PNG files are
        written, into a directory that exists, and an existing file is
        replaced only when the request says so (409 otherwise).
        """
        if not self.allowed():
            return
        query = self.query()
        try:
            if method == "GET":
                path = pathlib.Path(query.get("path", ""))
                kind = IMAGE_TYPES.get(path.suffix.lower())
                if not path.is_absolute() or kind is None:
                    return self.refuse(400, "not an image file")
                if path.stat().st_size > MAX_IMAGE_FILE_BYTES:
                    return self.refuse(413, "image too large")
                return self.send_body(path.read_bytes(), kind)
            directory = pathlib.Path(query.get("directory", ""))
            name = query.get("name", "")
            if not directory.is_absolute() or not directory.is_dir():
                return self.refuse(404, "no such directory")
            if not NAME.match(name) or not name.lower().endswith(".png"):
                return self.refuse(400, "the name must end in .png")
            length = int(self.headers.get("Content-Length", "-1"))
            if not 0 < length <= MAX_IMAGE_FILE_BYTES:
                return self.refuse(413, "picture too large")
            data = self.rfile.read(length)
            if not data.startswith(PNG_SIGNATURE):
                return self.refuse(400, "not a PNG picture")
            target = directory / name
            with notebook_lock:
                if target.exists() and query.get("replace") != "1":
                    return self.refuse(409, "exists")
                write_atomically(target, data)
            self.send_json({"saved": str(target)})
        except FileNotFoundError:
            self.refuse(404, "not found")
        except PermissionError:
            self.refuse(403, "not permitted there")
        except (OSError, ValueError) as error:
            self.refuse(500, str(error))

    @staticmethod
    def listing(directory, directories):
        """Names of the notebooks in the notebooks directory, or of the
        pages in a notebook.  A notebook that does not exist yet is empty."""
        if not directory.is_dir():
            return []
        return sorted(
            entry.name for entry in directory.iterdir()
            if NAME.match(entry.name) and
            (entry.is_dir() if directories else entry.is_file() and entry.name.endswith(".svg")))


def make_server(port, relay, notebooks, token, config, site=""):
    handler = type("ConfiguredHandler", (Handler,), {
        "relay": relay.rstrip("/"), "site": site_of(relay, site), "token": token,
        "notebooks": pathlib.Path(notebooks), "config": pathlib.Path(config)})
    handler = functools.partial(handler, directory=str(CLIENT_DIRECTORY))
    return http.server.ThreadingHTTPServer(("127.0.0.1", port), handler)


def main():
    parser = argparse.ArgumentParser(description="Kokkoja local program")
    parser.add_argument("--port", type=int, default=8760)
    parser.add_argument("--relay", default="ws://127.0.0.1:8770",
                        help="relay used for sessions started here "
                             "(default ws://127.0.0.1:8770)")
    parser.add_argument("--site", default="",
                        help="Kokkoja web site that invitations link to, so that "
                             "guests can join with nothing installed (default: "
                             "https://<host>/ when the relay is wss://<host>/relay)")
    parser.add_argument("--notebooks", type=pathlib.Path, default=None,
                        help="directory that holds the notebooks "
                             "(default: Kokkoja in your documents directory)")
    parser.add_argument("--config", type=pathlib.Path, default=None,
                        help="directory for Kokkoja's own settings "
                             "(default: your system's usual place)")
    parser.add_argument("--no-browser", action="store_true",
                        help="do not open a browser window")
    arguments = parser.parse_args()

    notebooks = arguments.notebooks or default_notebooks_directory()
    config = arguments.config or config_directory()
    token = load_token(config)
    server = make_server(arguments.port, arguments.relay, notebooks, token, config, arguments.site)
    address = f"http://localhost:{arguments.port}/#token={token}"
    print(f"Kokkoja is at {address}", flush=True)
    print(f"Notebooks are in {notebooks}", flush=True)
    print("Ctrl-C to stop.", flush=True)
    if not arguments.no_browser:
        webbrowser.open(address)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print()


if __name__ == "__main__":
    main()