#!/usr/bin/env python3
"""
Patch app.py: adds imports, /admin/upload-ota, /api/manifest, and upgrades
/app/version to report update_type ('native' | 'ota' | 'none').

Run from the same directory as app.py:
    python3 patch_ota_app.py
"""

path = "app.py"

with open(path, "r") as f:
    src = f.read()

# --- Patch 1: imports --------------------------------------------------------
old_1 = '''import base64
import time
import psutil
import requests as http_requests'''

new_1 = '''import base64
import time
import psutil
import requests as http_requests
import json
import uuid
import hashlib
import zipfile
import io
from datetime import datetime, timezone'''

assert old_1 in src, "Patch 1 (imports) anchor not found."
assert src.count(old_1) == 1, "Patch 1 anchor is not unique -- aborting."
src = src.replace(old_1, new_1)

# --- Patch 2: replace /app/version, add /admin/upload-ota + /api/manifest --
old_2 = '''@app.route("/app/version", methods=["GET"])
@require_auth
def app_version():
    """Pass the app's own version as ?current_version=1.2.0 and the backend
    tells you whether a newer one is available."""
    current_version = request.args.get("current_version")
    info = db.get_latest_app_version()
    if not info:
        return jsonify({"version": None, "apk_url": None, "notes": None, "update_available": False})

    info = dict(info)
    info["update_available"] = _version_is_newer(info.get("version"), current_version)
    return jsonify(info)'''

new_2 = '''def _content_type_for_ext(ext):
    ext = (ext or "").lower().lstrip(".")
    return {
        "png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
        "gif": "image/gif", "webp": "image/webp", "svg": "image/svg+xml",
        "ttf": "font/ttf", "otf": "font/otf", "woff": "font/woff", "woff2": "font/woff2",
        "json": "application/json",
    }.get(ext, "application/octet-stream")


def _asset_key(file_bytes):
    """Stable per-asset id (md5 of content) -- mirrors the 'key' field
    Expo's own export tooling uses for on-device asset caching."""
    return hashlib.md5(file_bytes).hexdigest()


def _sha256_base64url(file_bytes):
    digest = hashlib.sha256(file_bytes).digest()
    return base64.urlsafe_b64encode(digest).decode("utf-8").rstrip("=")


def _uuid_from_update_id(update_id):
    """Deterministic UUID derived from our own update_id, in the same spirit
    as Expo's reference server's convertSHA256HashToUUID()."""
    hex_digest = hashlib.sha256(update_id.encode("utf-8")).hexdigest()
    return str(uuid.UUID(hex_digest[:32]))


@app.route("/admin/upload-ota", methods=["POST"])
@require_admin
def admin_upload_ota():
    """Admin uploads a zip of the dist/ folder from
    `npx expo export --platform android --output-dir dist` (multipart
    'file' + form fields 'runtime_version', optional 'notes'). Every file
    inside gets pushed to the APP_UPDATES dataset repo under
    ota/<runtime_version>/<update_id>/..., then recorded as the latest OTA
    update -- same one-step upload-and-record flow as admin_upload_apk."""
    if "file" not in request.files:
        return jsonify({"error": "No file provided"}), 400

    file = request.files["file"]
    runtime_version = (request.form.get("runtime_version") or "").strip()
    notes = request.form.get("notes")

    if not runtime_version:
        return jsonify({"error": "runtime_version is required"}), 400

    update_id = str(uuid.uuid4())

    try:
        zf = zipfile.ZipFile(io.BytesIO(file.read()))
        names = zf.namelist()

        meta_name = next((n for n in names if n.endswith("metadata.json")), None)
        if not meta_name:
            return jsonify({"error": "metadata.json not found -- zip the dist/ folder produced by expo export"}), 400

        metadata = json.loads(zf.read(meta_name))
        base_dir = meta_name[: -len("metadata.json")]

        from huggingface_hub import HfApi
        api = HfApi(token=os.environ.get("HF_TOKEN") or os.environ.get("Hf_token"))

        def upload_asset(rel_path):
            data = zf.read(base_dir + rel_path)
            path_in_repo = f"ota/{runtime_version}/{update_id}/{rel_path}"
            api.upload_file(
                path_or_fileobj=data,
                path_in_repo=path_in_repo,
                repo_id=APP_UPDATE_DATASET_REPO,
                repo_type="dataset",
            )
            url = f"https://huggingface.co/datasets/{APP_UPDATE_DATASET_REPO}/resolve/main/{path_in_repo}"
            return url, data

        android_meta = metadata["fileMetadata"]["android"]

        bundle_url, bundle_bytes = upload_asset(android_meta["bundle"])
        launch_asset = {
            "hash": _sha256_base64url(bundle_bytes),
            "key": _asset_key(bundle_bytes),
            "contentType": "application/javascript",
            "url": bundle_url,
        }

        assets = []
        for a in android_meta.get("assets", []):
            url, data = upload_asset(a["path"])
            assets.append({
                "hash": _sha256_base64url(data),
                "key": _asset_key(data),
                "contentType": _content_type_for_ext(a.get("ext")),
                "url": url,
            })

        db.set_latest_ota_update(update_id, runtime_version, json.dumps(launch_asset), json.dumps(assets), notes)

    except KeyError as e:
        return jsonify({"error": f"Unexpected metadata.json shape, missing {e}"}), 400
    except Exception as e:
        return jsonify({"error": f"OTA upload failed: {e}"}), 500

    return jsonify({
        "status": "updated",
        "update_id": update_id,
        "runtime_version": runtime_version,
        "asset_count": len(assets) + 1,
    })


@app.route("/api/manifest", methods=["GET"])
def api_manifest():
    """The endpoint expo-updates itself calls (set as `updates.url` in
    app.json) -- separate from /app/version below, which only powers the
    manual 'Check for updates' button's native-vs-OTA decision. Implements
    the Expo Updates protocol's plain-JSON response path (no code signing;
    protocol version 1's 'no update' case is a 406, per spec)."""
    protocol_version = request.headers.get("expo-protocol-version", "0")
    platform = request.headers.get("expo-platform") or request.args.get("platform")
    runtime_version = request.headers.get("expo-runtime-version") or request.args.get("runtime-version")

    if platform != "android":
        return jsonify({"error": "Unsupported platform. Expected android."}), 400
    if not runtime_version:
        return jsonify({"error": "No runtimeVersion provided."}), 400

    update = db.get_latest_ota_update()
    if not update or update.get("runtime_version") != runtime_version:
        return jsonify({"error": "No update available for this runtime version."}), 404

    manifest_id = _uuid_from_update_id(update["update_id"])
    current_update_id = request.headers.get("expo-current-update-id")

    if protocol_version == "1" and current_update_id == manifest_id:
        return jsonify({"error": "No update available."}), 406

    created_at = datetime.fromtimestamp(update["created_at"], tz=timezone.utc).isoformat()

    manifest = {
        "id": manifest_id,
        "createdAt": created_at,
        "runtimeVersion": runtime_version,
        "launchAsset": json.loads(update["launch_asset_json"]),
        "assets": json.loads(update["assets_json"]),
        "metadata": {},
        "extra": {},
    }

    resp = jsonify(manifest)
    resp.headers["expo-protocol-version"] = protocol_version
    resp.headers["expo-sfv-version"] = "0"
    resp.headers["cache-control"] = "private, max-age=0"
    return resp


@app.route("/app/version", methods=["GET"])
@require_auth
def app_version():
    """Pass ?current_version=1.2.0 (native app version) and
    ?current_runtime_version=... (expo-updates runtime version) and the
    backend tells you which kind of update -- if any -- is waiting:
    'native' (needs the full-APK flow) or 'ota' (safe to hot-load via
    expo-updates' own checkForUpdateAsync, which talks to /api/manifest)."""
    current_version = request.args.get("current_version")
    current_runtime_version = request.args.get("current_runtime_version")

    native_info = db.get_latest_app_version()
    if native_info and _version_is_newer(native_info.get("version"), current_version):
        info = dict(native_info)
        info["update_available"] = True
        info["update_type"] = "native"
        return jsonify(info)

    ota_info = db.get_latest_ota_update()
    if ota_info and current_runtime_version and ota_info.get("runtime_version") == current_runtime_version:
        return jsonify({
            "version": None,
            "apk_url": None,
            "notes": ota_info.get("notes"),
            "update_available": True,
            "update_type": "ota",
        })

    return jsonify({"version": None, "apk_url": None, "notes": None, "update_available": False, "update_type": "none"})'''

assert old_2 in src, "Patch 2 (/app/version + new routes) anchor not found."
assert src.count(old_2) == 1, "Patch 2 anchor is not unique -- aborting."
src = src.replace(old_2, new_2)

with open(path, "w") as f:
    f.write(src)

print("app.py patched successfully.")
print("  - added _content_type_for_ext, _asset_key, _sha256_base64url, _uuid_from_update_id")
print("  - added POST /admin/upload-ota")
print("  - added GET  /api/manifest")
print("  - updated GET /app/version (now returns update_type)")