#!/usr/bin/env python3
"""JARVIS update server with a narrow, authenticated publisher API.

The future Windows JARVIS Asset app stages a release and then explicitly publishes it.
Only the update namespace is writable; manifest.json is moved last so Android clients never
observe a partial release.
"""

from __future__ import annotations

import argparse
import hashlib
import hmac
import json
import os
import re
import secrets
import shutil
import tempfile
from email.utils import formatdate
from http import HTTPStatus
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import unquote, urlparse


SAFE_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,159}$")
SAFE_RELEASE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,79}$")
ALLOWED_SUFFIXES = {".apk", ".json", ".md", ".txt", ".zip"}
MAX_BYTES = 700 * 1024 * 1024
MAX_TEXT_BYTES = 2 * 1024 * 1024
REQUIRED = {"manifest.json", "release-notes.md", "changes.json"}
PERSONAL_PACKAGE = "it.jarvis.app.personal.ecosystem"
RANGE_HEADER = re.compile(r"^bytes=(\d*)-(\d*)$")


class PublisherServer(ThreadingHTTPServer):
    def __init__(self, address, handler, root: Path, token: str):
        super().__init__(address, handler)
        self.root = root.resolve()
        # Android clients request /jarvis-ecosystem/manifest.json from this static root.
        self.live = (self.root / "jarvis-ecosystem").resolve()
        self.staging = (self.root / ".staging").resolve()
        self.token = token
        self.live.mkdir(parents=True, exist_ok=True)
        self.staging.mkdir(parents=True, exist_ok=True)


class Handler(SimpleHTTPRequestHandler):
    server: PublisherServer
    range_remaining = None

    def guess_type(self, path):
        if path.lower().endswith(".apk"):
            return "application/vnd.android.package-archive"
        return super().guess_type(path)

    def end_headers(self):
        request_path = urlparse(self.path).path
        local_path = Path(self.translate_path(self.path))
        if local_path.is_file():
            self.send_header("ETag", self._file_etag(local_path))
        if request_path.endswith("/manifest.json") or request_path.endswith("/catalog.json") or request_path.lower().endswith(".zip"):
            self.send_header("Cache-Control", "no-store, no-cache, must-revalidate, max-age=0")
            self.send_header("Pragma", "no-cache")
            self.send_header("Expires", "0")
        if request_path.lower().endswith(".apk"):
            filename = Path(unquote(request_path)).name
            self.send_header("Content-Disposition", f'attachment; filename="{filename}"')
            self.send_header("X-Content-Type-Options", "nosniff")
        self.send_header("Accept-Ranges", "bytes")
        super().end_headers()

    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=str(args[2].root) if len(args) > 2 else None, **kwargs)

    def log_message(self, fmt, *args):
        print(f"{self.client_address[0]} - {fmt % args}", flush=True)

    def do_GET(self):
        if urlparse(self.path).path == "/health":
            return self._json(HTTPStatus.OK, {"ok": True, "service": "jarvis-update-server", "publisherApi": 1})
        if self.headers.get("Range") and self._if_range_matches():
            source = self._send_range_headers()
            if source is not None:
                try:
                    self.copyfile(source, self.wfile)
                finally:
                    source.close()
            return
        return super().do_GET()

    def do_HEAD(self):
        if self.headers.get("Range") and self._if_range_matches():
            source = self._send_range_headers()
            if source is not None:
                source.close()
            return
        return super().do_HEAD()

    @staticmethod
    def _file_etag(path: Path):
        stat = path.stat()
        return f'"{stat.st_size:x}-{stat.st_mtime_ns:x}"'

    def _if_range_matches(self):
        condition = self.headers.get("If-Range", "").strip()
        if not condition:
            return True
        path = Path(self.translate_path(self.path))
        if not path.is_file():
            return False
        # Android DownloadManager uses the entity validator to decide whether its partial bytes
        # still belong to this exact APK. A mismatch correctly falls back to a full 200 response.
        if condition.startswith('"'):
            return hmac.compare_digest(condition, self._file_etag(path))
        return condition == formatdate(path.stat().st_mtime, usegmt=True)

    def _send_range_headers(self):
        path = Path(self.translate_path(self.path))
        if not path.is_file():
            self.send_error(HTTPStatus.NOT_FOUND, "File not found")
            return None
        size = path.stat().st_size
        match = RANGE_HEADER.fullmatch(self.headers.get("Range", "").strip())
        selected = None
        if match and size > 0:
            first, last = match.groups()
            if not first:
                suffix = int(last or "0")
                if suffix > 0:
                    selected = (max(0, size - suffix), size - 1)
            else:
                start = int(first)
                end = int(last) if last else size - 1
                if start < size and end >= start:
                    selected = (start, min(end, size - 1))
        if selected is None:
            self.send_response(HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE)
            self.send_header("Content-Range", f"bytes */{size}")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return None
        start, end = selected
        source = path.open("rb")
        source.seek(start)
        self.range_remaining = end - start + 1
        self.send_response(HTTPStatus.PARTIAL_CONTENT)
        self.send_header("Content-Type", self.guess_type(str(path)))
        self.send_header("Content-Range", f"bytes {start}-{end}/{size}")
        self.send_header("Content-Length", str(self.range_remaining))
        self.send_header("Last-Modified", formatdate(path.stat().st_mtime, usegmt=True))
        self.end_headers()
        return source

    def copyfile(self, source, outputfile):
        remaining = self.range_remaining
        if remaining is None:
            return super().copyfile(source, outputfile)
        try:
            while remaining > 0:
                chunk = source.read(min(256 * 1024, remaining))
                if not chunk:
                    break
                outputfile.write(chunk)
                remaining -= len(chunk)
        finally:
            self.range_remaining = None

    def do_PUT(self):
        if not self._authorized():
            return self._json(HTTPStatus.UNAUTHORIZED, {"error": "unauthorized"})
        parts = self._parts("/publisher/v1/stage/")
        if parts is None or len(parts) != 2:
            return self._json(HTTPStatus.NOT_FOUND, {"error": "unknown_endpoint"})
        release_id, filename = parts
        if not SAFE_RELEASE.fullmatch(release_id) or not self._valid_filename(filename):
            return self._json(HTTPStatus.BAD_REQUEST, {"error": "invalid_name"})
        length = self._content_length()
        limit = MAX_BYTES if filename.lower().endswith((".apk", ".zip")) else MAX_TEXT_BYTES
        if length is None or length < 1 or length > limit:
            return self._json(HTTPStatus.REQUEST_ENTITY_TOO_LARGE, {"error": "invalid_size", "limit": limit})
        folder = (self.server.staging / release_id).resolve()
        if folder.parent != self.server.staging:
            return self._json(HTTPStatus.BAD_REQUEST, {"error": "invalid_release"})
        folder.mkdir(parents=True, exist_ok=True)
        destination = (folder / filename).resolve()
        if destination.parent != folder:
            return self._json(HTTPStatus.BAD_REQUEST, {"error": "invalid_path"})
        with tempfile.NamedTemporaryFile(dir=folder, delete=False) as temp:
            remaining = length
            while remaining:
                chunk = self.rfile.read(min(1024 * 1024, remaining))
                if not chunk:
                    Path(temp.name).unlink(missing_ok=True)
                    return self._json(HTTPStatus.BAD_REQUEST, {"error": "incomplete_body"})
                temp.write(chunk)
                remaining -= len(chunk)
            temp_path = Path(temp.name)
        os.replace(temp_path, destination)
        return self._json(HTTPStatus.CREATED, {
            "ok": True,
            "releaseId": release_id,
            "file": filename,
            "bytes": length,
            "sha256": self._sha256(destination),
            "state": "staged",
        })

    def do_POST(self):
        if not self._authorized():
            return self._json(HTTPStatus.UNAUTHORIZED, {"error": "unauthorized"})
        parts = self._parts("/publisher/v1/publish/")
        if parts is None or len(parts) != 1 or not SAFE_RELEASE.fullmatch(parts[0]):
            return self._json(HTTPStatus.NOT_FOUND, {"error": "unknown_endpoint"})
        release_id = parts[0]
        folder = (self.server.staging / release_id).resolve()
        if folder.parent != self.server.staging or not folder.is_dir():
            return self._json(HTTPStatus.NOT_FOUND, {"error": "release_not_staged"})
        try:
            manifest, apk_name = self._validate_release(folder)
            for item in folder.iterdir():
                if item.is_file() and item.name != "manifest.json":
                    os.replace(item, self.server.live / item.name)
            os.replace(folder / "manifest.json", self.server.live / "manifest.json")
            shutil.rmtree(folder, ignore_errors=True)
            return self._json(HTTPStatus.OK, {
                "ok": True,
                "state": "published",
                "releaseId": release_id,
                "versionCode": manifest["android"]["versionCode"],
                "versionName": manifest["android"]["versionName"],
                "apk": apk_name,
            })
        except (ValueError, KeyError, json.JSONDecodeError) as exc:
            return self._json(HTTPStatus.BAD_REQUEST, {"error": "invalid_release", "detail": str(exc)})

    def _validate_release(self, folder: Path):
        names = {item.name for item in folder.iterdir() if item.is_file()}
        missing = REQUIRED - names
        if missing:
            raise ValueError("missing: " + ", ".join(sorted(missing)))
        manifest = json.loads((folder / "manifest.json").read_text(encoding="utf-8"))
        android = manifest["android"]
        if android.get("packageName") != PERSONAL_PACKAGE:
            raise ValueError("unexpected packageName")
        if not isinstance(android.get("versionCode"), int) or android["versionCode"] < 1:
            raise ValueError("invalid versionCode")
        current_manifest_path = self.server.live / "manifest.json"
        if current_manifest_path.is_file():
            current = json.loads(current_manifest_path.read_text(encoding="utf-8"))
            current_android = current.get("android", {})
            if current_android.get("packageName") == PERSONAL_PACKAGE:
                current_code = current_android.get("versionCode", 0)
                if not isinstance(current_code, int) or android["versionCode"] <= current_code:
                    raise ValueError("versionCode must be greater than the currently published Personal build")
        apk_name = android.get("url", "")
        if not self._valid_filename(apk_name) or not apk_name.lower().endswith(".apk"):
            raise ValueError("invalid APK filename")
        apk = folder / apk_name
        if not apk.is_file():
            raise ValueError("APK not staged")
        expected = str(android.get("sha256", "")).lower()
        if not re.fullmatch(r"[0-9a-f]{64}", expected) or not hmac.compare_digest(expected, self._sha256(apk)):
            raise ValueError("APK sha256 mismatch")
        if android.get("summaryUrl") != "release-notes.md" or android.get("detailsUrl") != "changes.json":
            raise ValueError("release document paths must be local canonical names")
        json.loads((folder / "changes.json").read_text(encoding="utf-8"))
        if not (folder / "release-notes.md").read_text(encoding="utf-8").strip():
            raise ValueError("empty release notes")
        return manifest, apk_name

    def _authorized(self):
        header = self.headers.get("Authorization", "")
        prefix = "Bearer "
        return header.startswith(prefix) and hmac.compare_digest(header[len(prefix):], self.server.token)

    def _parts(self, prefix):
        path = unquote(urlparse(self.path).path)
        if not path.startswith(prefix):
            return None
        return [part for part in path[len(prefix):].split("/") if part]

    def _valid_filename(self, name):
        return bool(SAFE_NAME.fullmatch(name)) and Path(name).suffix.lower() in ALLOWED_SUFFIXES

    def _content_length(self):
        try:
            return int(self.headers.get("Content-Length", ""))
        except ValueError:
            return None

    @staticmethod
    def _sha256(path: Path):
        digest = hashlib.sha256()
        with path.open("rb") as stream:
            for chunk in iter(lambda: stream.read(1024 * 1024), b""):
                digest.update(chunk)
        return digest.hexdigest()

    def _json(self, status, payload):
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)


def token_from(path: Path) -> str:
    path.parent.mkdir(parents=True, exist_ok=True)
    if not path.exists():
        path.write_text(secrets.token_urlsafe(48), encoding="utf-8")
    token = path.read_text(encoding="utf-8").strip()
    if len(token) < 32:
        raise SystemExit("publisher token must contain at least 32 characters")
    return token


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8795)
    parser.add_argument("--token-file", type=Path, default=Path(os.environ.get(
        "LOCALAPPDATA", str(Path.home()))) / "JARVIS" / "update-server" / "publisher.token")
    args = parser.parse_args()
    server = PublisherServer((args.host, args.port), Handler, args.root, token_from(args.token_file))
    print(f"JARVIS update server on http://{args.host}:{args.port}; root={server.root}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    main()