#!/usr/bin/env python3
"""Pull Navi's encrypted history snapshot, decrypt it, leave a gzip-JSONL file.

Runs on the consumer side — a machine that holds the private half of the age
key and never shares it. Dependencies: Python 3.11+ and `pip install pyrage`.

    # once, on the consumer: make a keypair, send the *public* half to Navi
    ./navi_export.py keygen --out ~/.config/navi-export
    # -> writes key.txt (secret, mode 600) and prints age1... for NAVI_EXPORT_AGE_RECIPIENT

    # every run (cron-friendly)
    ./navi_export.py pull \
        --url https://navi.example \
        --token-file ~/.config/navi-export/token \
        --identity ~/.config/navi-export/key.txt \
        --out ~/ngram/messages.jsonl.gz

Exit codes: 0 updated, 3 unchanged (nothing downloaded), 1 error.
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import os
import sys
import urllib.error
import urllib.request
from pathlib import Path

MANIFEST_PATH = "/export/v1/manifest"
MESSAGES_PATH = "/export/v1/messages"
USER_AGENT = "navi-export/1"


class PullError(Exception):
    pass


def _die(message: str) -> int:
    print(f"error: {message}", file=sys.stderr)
    return 1


def _request(url: str, token: str, *, accept: str = "application/json") -> urllib.request.Request:
    return urllib.request.Request(
        url,
        headers={
            "Authorization": f"Bearer {token}",
            "Accept": accept,
            "User-Agent": USER_AGENT,
        },
    )


def fetch_manifest(base: str, token: str, timeout: float) -> dict:
    try:
        with urllib.request.urlopen(_request(base + MANIFEST_PATH, token), timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        if exc.code == 401:
            raise PullError("the server rejected the token (401)") from exc
        if exc.code == 503:
            raise PullError("the server has no export built yet (503)") from exc
        raise PullError(f"manifest request failed: HTTP {exc.code}") from exc
    except urllib.error.URLError as exc:
        raise PullError(f"cannot reach {base}: {exc.reason}") from exc


def _read_token(args) -> str:
    if args.token:
        return args.token.strip()
    if args.token_file:
        path = Path(args.token_file).expanduser()
        try:
            return path.read_text(encoding="utf-8").strip()
        except OSError as exc:
            raise PullError(f"cannot read token file {path}: {exc}") from exc
    raise PullError("one of --token or --token-file is required")


def _load_state(path: Path) -> dict:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return {}


def _save_state(path: Path, state: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(state, indent=2) + "\n", encoding="utf-8")
    os.replace(tmp, path)


def _download(base: str, token: str, timeout: float, dest: Path) -> str:
    """Stream the ciphertext to `dest`; return the sha256 of what landed."""
    digest = hashlib.sha256()
    try:
        with urllib.request.urlopen(
            _request(base + MESSAGES_PATH, token, accept="application/octet-stream"),
            timeout=timeout,
        ) as resp, open(dest, "wb") as out:
            while True:
                block = resp.read(1 << 20)
                if not block:
                    break
                digest.update(block)
                out.write(block)
    except urllib.error.HTTPError as exc:
        dest.unlink(missing_ok=True)
        raise PullError(f"download failed: HTTP {exc.code}") from exc
    except urllib.error.URLError as exc:
        dest.unlink(missing_ok=True)
        raise PullError(f"download failed: {exc.reason}") from exc
    return digest.hexdigest()


def _decrypt(identity_path: Path, src: Path, dest: Path) -> None:
    import pyrage

    try:
        identity = pyrage.x25519.Identity.from_str(
            identity_path.expanduser().read_text(encoding="utf-8").strip()
        )
    except (OSError, pyrage.IdentityError) as exc:
        raise PullError(f"cannot use identity {identity_path}: {exc}") from exc

    try:
        with open(src, "rb") as enc, open(dest, "wb") as out:
            pyrage.decrypt_io(enc, out, [identity])
    except pyrage.DecryptError as exc:
        dest.unlink(missing_ok=True)
        raise PullError(
            "decryption failed — the file is not for this identity, or it is corrupt"
        ) from exc


def _verify_gzip(path: Path) -> int:
    """Fully decompress once: catches truncation now, not at n-gram build time."""
    total = 0
    try:
        with gzip.open(path, "rb") as f:
            for block in iter(lambda: f.read(1 << 20), b""):
                total += len(block)
    except (OSError, EOFError) as exc:
        raise PullError(f"decrypted file is not a valid gzip stream: {exc}") from exc
    return total


def cmd_pull(args) -> int:
    base = args.url.rstrip("/")
    out = Path(args.out).expanduser()
    state_path = Path(args.state).expanduser() if args.state else out.with_suffix(out.suffix + ".state.json")
    token = _read_token(args)
    state = _load_state(state_path)

    manifest = fetch_manifest(base, token, args.timeout)
    content_id = manifest.get("content_id", "")
    if not content_id or not manifest.get("file", {}).get("sha256"):
        raise PullError("manifest is missing content_id/sha256 — server and script disagree")

    if not args.force and state.get("content_id") == content_id and out.exists():
        print(f"unchanged since {manifest.get('built_at')} (content_id {content_id[:12]})")
        return 3

    out.parent.mkdir(parents=True, exist_ok=True)
    part_enc = out.with_name(out.name + ".part.age")
    part_plain = out.with_name(out.name + ".part")
    try:
        got = _download(base, token, args.timeout, part_enc)
        expected = manifest["file"]["sha256"]
        if got != expected:
            raise PullError(f"checksum mismatch: got {got[:12]}, manifest says {expected[:12]}")

        _decrypt(Path(args.identity), part_enc, part_plain)
        plain_bytes = _verify_gzip(part_plain)
        os.replace(part_plain, out)
    finally:
        part_enc.unlink(missing_ok=True)
        part_plain.unlink(missing_ok=True)

    _save_state(state_path, {"content_id": content_id, "built_at": manifest.get("built_at")})

    window = manifest.get("window", {})
    span = f"{window.get('from') or '?'} .. {window.get('to') or '?'}"
    if window.get("truncated"):
        span += f" (window capped at {window.get('max_bytes', 0) / 1024 / 1024:.0f} MB)"
    print(
        f"updated {out}: {plain_bytes / 1024 / 1024:.1f} MB, "
        f"{manifest.get('rows', 0)} rows, {manifest.get('sessions', 0)} sessions, {span}"
    )
    print(f"built_at {manifest.get('built_at')}  content_id {content_id[:12]}")
    return 0


def cmd_keygen(args) -> int:
    import pyrage

    out = Path(args.out).expanduser()
    out.mkdir(parents=True, exist_ok=True)
    key_path = out / "key.txt"
    if key_path.exists() and not args.force:
        return _die(f"{key_path} already exists (use --force to overwrite — it invalidates the old export)")
    identity = pyrage.x25519.Identity.generate()
    key_path.write_text(str(identity) + "\n", encoding="utf-8")
    os.chmod(key_path, 0o600)
    print(key_path)
    print(str(identity.to_public()))
    print("Set NAVI_EXPORT_AGE_RECIPIENT on Navi to the line above.", file=sys.stderr)
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(prog="navi-export", description=__doc__.splitlines()[0])
    sub = p.add_subparsers(dest="cmd", required=True)

    pull = sub.add_parser("pull", help="fetch, decrypt and install the snapshot")
    pull.add_argument("--url", required=True, help="Navi base URL, e.g. https://navi.example")
    pull.add_argument("--token", help="export token (prefer --token-file)")
    pull.add_argument("--token-file", help="file holding the export token")
    pull.add_argument("--identity", required=True, help="age secret key file")
    pull.add_argument("--out", required=True, help="where to write the decrypted gzip-JSONL")
    pull.add_argument("--state", help="state file (default: <out>.state.json)")
    pull.add_argument("--timeout", type=float, default=120.0, help="per-request timeout, seconds")
    pull.add_argument("--force", action="store_true", help="download even if content_id is unchanged")
    pull.set_defaults(func=cmd_pull)

    keygen = sub.add_parser("keygen", help="generate an age keypair for the consumer")
    keygen.add_argument("--out", default=str(Path.home() / ".config" / "navi-export"))
    keygen.add_argument("--force", action="store_true")
    keygen.set_defaults(func=cmd_keygen)
    return p


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        return args.func(args)
    except PullError as exc:
        return _die(str(exc))
    except KeyboardInterrupt:
        return 130


if __name__ == "__main__":
    sys.exit(main())