#!/usr/bin/env python3
"""Audit a built Solarma quarterly place-pack release before publication."""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import math
import re
import sys
from pathlib import Path
from typing import Any, Iterable


GEOHASH_ALPHABET = "0123456789bcdefghjkmnpqrstuvwxyz"
GEOHASH_BITS = (16, 8, 4, 2, 1)
PACK_KEY_RE = re.compile(r"^[0-9bcdefghjkmnpqrstuvwxyz]{3}$")
SHA256_RE = re.compile(r"^[0-9a-f]{64}$")
FORBIDDEN_KEYS = {"uid", "user", "username", "changeset", "version", "timestamp"}
ACTIVITY_IDS = {
    "playground",
    "park",
    "walking",
    "hiking",
    "discGolf",
    "dogPark",
    "bike",
    "court",
    "skate",
    "picnic",
    "water",
    "fishing",
    "viewpoint",
    "birding",
    "golf",
}
EXPECTED_PACK_KEYS = {
    "schemaVersion",
    "datasetVersion",
    "quarter",
    "sourceTimestamp",
    "geohash",
    "attribution",
    "places",
}
EXPECTED_PLACE_KEYS = {
    "access",
    "activityIds",
    "address",
    "fee",
    "id",
    "latitude",
    "lit",
    "litTag",
    "longitude",
    "name",
    "osm",
    "tags",
}
REPRESENTATIVE_COVERAGE = {
    "Tucson": (32.2226, -110.9747),
    "New York City": (40.7128, -74.0060),
    "Washington, DC": (38.9072, -77.0369),
    "Anchorage": (61.2181, -149.9003),
    "Honolulu": (21.3099, -157.8581),
    "San Juan": (18.4655, -66.1057),
    "Charlotte Amalie": (18.3419, -64.9307),
}


class AuditError(RuntimeError):
    pass


def canonical_json_bytes(value: Any) -> bytes:
    return (
        json.dumps(
            value,
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
            allow_nan=False,
        ).encode("utf-8")
        + b"\n"
    )


def sha256_file(path: Path) -> str:
    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 encode_geohash(latitude: float, longitude: float, precision: int = 3) -> str:
    lat_range = [-90.0, 90.0]
    lon_range = [-180.0, 180.0]
    even = True
    bit_index = 0
    current = 0
    encoded: list[str] = []
    while len(encoded) < precision:
        bounds = lon_range if even else lat_range
        value = longitude if even else latitude
        midpoint = (bounds[0] + bounds[1]) / 2
        if value >= midpoint:
            current |= GEOHASH_BITS[bit_index]
            bounds[0] = midpoint
        else:
            bounds[1] = midpoint
        even = not even
        if bit_index < 4:
            bit_index += 1
        else:
            encoded.append(GEOHASH_ALPHABET[current])
            bit_index = 0
            current = 0
    return "".join(encoded)


def haversine_km(a: tuple[float, float], b: tuple[float, float]) -> float:
    radius = 6371.0
    lat1, lon1 = map(math.radians, a)
    lat2, lon2 = map(math.radians, b)
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    value = (
        math.sin(dlat / 2) ** 2
        + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2
    )
    return 2 * radius * math.asin(min(1.0, math.sqrt(value)))


def walk_keys(value: Any) -> Iterable[str]:
    if isinstance(value, dict):
        for key, child in value.items():
            yield str(key)
            yield from walk_keys(child)
    elif isinstance(value, list):
        for child in value:
            yield from walk_keys(child)


def require(condition: bool, message: str) -> None:
    if not condition:
        raise AuditError(message)


def audit_sources(source_manifest_path: Path, output_sources: list[dict[str, Any]]) -> None:
    source_manifest = json.loads(source_manifest_path.read_text(encoding="utf-8"))
    extracts = source_manifest["extracts"]
    expected = sorted(
        (
            {
                "fileName": Path(item["path"]).name,
                "url": item["url"],
                "sha256": item["sha256"],
            }
            for item in extracts
        ),
        key=lambda item: (item["url"], item["sha256"]),
    )
    require(output_sources == expected, "output sources differ from the source manifest")
    for item in extracts:
        source_path = (source_manifest_path.parent / item["path"]).resolve()
        require(source_path.is_file(), f"missing source extract: {source_path}")
        require(
            sha256_file(source_path) == item["sha256"],
            f"source SHA-256 mismatch: {source_path.name}",
        )


def audit_release(release_dir: Path, source_manifest_path: Path | None) -> dict[str, Any]:
    release_dir = release_dir.resolve()
    manifest_path = release_dir / "manifest.json"
    attribution_path = release_dir / "attribution.json"
    headers_path = release_dir / "deploy-headers.json"
    require(manifest_path.is_file(), "manifest.json is missing")
    require(attribution_path.is_file(), "attribution.json is missing")
    require(headers_path.is_file(), "deploy-headers.json is missing")

    manifest_raw = manifest_path.read_bytes()
    manifest = json.loads(manifest_raw)
    attribution_raw = attribution_path.read_bytes()
    attribution = json.loads(attribution_raw)
    headers = json.loads(headers_path.read_bytes())
    require(manifest_raw == canonical_json_bytes(manifest), "manifest is not canonical JSON")
    require(attribution_raw == canonical_json_bytes(attribution), "attribution is not canonical JSON")
    require(manifest["schemaVersion"] == 1, "unexpected schemaVersion")
    require(manifest["dataset"] == "solarma-osm-outing-places", "unexpected dataset")
    require(manifest["quarter"] == "2026-Q3", "unexpected quarter")
    require(manifest["sourceTimestamp"] == "2026-08-30T20:21:06Z", "unexpected source timestamp")
    require(manifest["geohashPrecision"] == 3, "unexpected geohash precision")
    require(manifest["cacheMaxAgeSeconds"] == 90 * 24 * 60 * 60, "unexpected cache age")
    require(manifest["attribution"] == attribution, "attribution files disagree")

    descriptors: dict[str, dict[str, Any]] = manifest["packs"]
    actual_files = {
        path.relative_to(release_dir).as_posix()
        for path in (release_dir / "packs").glob("*.json.gz")
    }
    described_files = {item["path"] for item in descriptors.values()}
    require(actual_files == described_files, "pack file set differs from manifest")

    seen_ids: set[str] = set()
    place_count = 0
    compressed_bytes = 0
    uncompressed_bytes = 0
    coverage: dict[str, dict[str, Any]] = {}
    coverage_by_hash: dict[str, list[tuple[str, tuple[float, float]]]] = {}
    for name, point in REPRESENTATIVE_COVERAGE.items():
        coverage_by_hash.setdefault(encode_geohash(*point), []).append((name, point))

    for geohash, descriptor in sorted(descriptors.items()):
        require(bool(PACK_KEY_RE.fullmatch(geohash)), f"invalid geohash key: {geohash}")
        relative = f"packs/{geohash}.json.gz"
        require(descriptor["path"] == relative, f"path mismatch for {geohash}")
        require(descriptor["contentEncoding"] == "gzip", f"encoding mismatch for {geohash}")
        require(
            descriptor["contentType"] == "application/json; charset=utf-8",
            f"content type mismatch for {geohash}",
        )
        require(bool(SHA256_RE.fullmatch(descriptor["sha256"])), f"bad SHA-256 for {geohash}")
        require(
            descriptor["etag"] == f'"sha256-{descriptor["sha256"]}"',
            f"ETag mismatch for {geohash}",
        )

        path = release_dir / relative
        compressed = path.read_bytes()
        require(len(compressed) == descriptor["compressedBytes"], f"compressed size mismatch: {relative}")
        require(hashlib.sha256(compressed).hexdigest() == descriptor["sha256"], f"SHA mismatch: {relative}")
        require(
            compressed[:10] == b"\x1f\x8b\x08\x00\x00\x00\x00\x00\x02\xff",
            f"non-deterministic gzip header: {relative}",
        )
        raw = gzip.decompress(compressed)
        require(len(raw) == descriptor["uncompressedBytes"], f"uncompressed size mismatch: {relative}")
        pack = json.loads(raw)
        require(raw == canonical_json_bytes(pack), f"non-canonical pack JSON: {relative}")
        require(set(pack) == EXPECTED_PACK_KEYS, f"unexpected pack keys: {relative}")
        require(pack["schemaVersion"] == manifest["schemaVersion"], f"schema mismatch: {relative}")
        require(pack["datasetVersion"] == manifest["datasetVersion"], f"dataset mismatch: {relative}")
        require(pack["quarter"] == manifest["quarter"], f"quarter mismatch: {relative}")
        require(pack["sourceTimestamp"] == manifest["sourceTimestamp"], f"timestamp mismatch: {relative}")
        require(pack["geohash"] == geohash, f"geohash mismatch: {relative}")
        require(pack["attribution"] == attribution, f"attribution mismatch: {relative}")
        require(len(pack["places"]) == descriptor["count"], f"place count mismatch: {relative}")

        forbidden = {key.lower() for key in walk_keys(pack)} & FORBIDDEN_KEYS
        require(not forbidden, f"private OSM metadata keys in {relative}: {sorted(forbidden)}")

        for place in pack["places"]:
            require(set(place) == EXPECTED_PLACE_KEYS, f"unexpected place keys in {relative}")
            require(isinstance(place["id"], str) and place["id"], f"invalid place id in {relative}")
            require(place["id"] not in seen_ids, f"duplicate place id: {place['id']}")
            seen_ids.add(place["id"])
            latitude = place["latitude"]
            longitude = place["longitude"]
            require(isinstance(latitude, (int, float)) and math.isfinite(latitude), f"bad latitude: {place['id']}")
            require(isinstance(longitude, (int, float)) and math.isfinite(longitude), f"bad longitude: {place['id']}")
            require(-90 <= latitude <= 90 and -180 <= longitude <= 180, f"coordinate out of range: {place['id']}")
            require(encode_geohash(latitude, longitude) == geohash, f"place in wrong pack: {place['id']}")
            require(
                isinstance(place["activityIds"], list)
                and bool(place["activityIds"])
                and set(place["activityIds"]) <= ACTIVITY_IDS,
                f"invalid activities: {place['id']}",
            )

        for name, point in coverage_by_hash.get(geohash, []):
            nearest = min(
                (haversine_km(point, (item["latitude"], item["longitude"])), item["id"])
                for item in pack["places"]
            )
            require(nearest[0] <= 100, f"no representative place within 100 km of {name}")
            coverage[name] = {
                "geohash": geohash,
                "nearestPlaceId": nearest[1],
                "nearestKm": round(nearest[0], 2),
            }

        place_count += len(pack["places"])
        compressed_bytes += len(compressed)
        uncompressed_bytes += len(raw)

    require(place_count == manifest["placeCount"], "manifest placeCount differs from pack sum")
    require(set(coverage) == set(REPRESENTATIVE_COVERAGE), "representative coverage is incomplete")
    require(
        set(headers["headers"]) >= {"manifest.json", "attribution.json", *described_files},
        "deployment headers are incomplete",
    )
    if source_manifest_path is not None:
        audit_sources(source_manifest_path.resolve(), manifest["sources"])

    return {
        "datasetVersion": manifest["datasetVersion"],
        "sourceTimestamp": manifest["sourceTimestamp"],
        "packCount": len(descriptors),
        "placeCount": place_count,
        "compressedBytes": compressed_bytes,
        "uncompressedBytes": uncompressed_bytes,
        "uniquePlaceIds": len(seen_ids),
        "representativeCoverage": coverage,
        "sourceExtractsRehashed": len(manifest["sources"]) if source_manifest_path else 0,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("release_dir", type=Path)
    parser.add_argument("--source-manifest", type=Path)
    args = parser.parse_args()
    try:
        summary = audit_release(args.release_dir, args.source_manifest)
    except (AuditError, KeyError, OSError, ValueError, json.JSONDecodeError) as error:
        print(f"audit failed: {error}", file=sys.stderr)
        return 1
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())