scripts/scrapers/pbf_processor.py

"""
pbf_processor.py
================
High-speed OSM .pbf file processor using pyosmium.
Extracts restaurants, hotels, attractions, and cultural sites with
full tag data (cuisine, dietary, opening hours, cultural info, etc.)

Processing speeds (approx):
  asia.osm.pbf   (15 GB) → ~45–90 min on modern hardware
  australia.pbf  ( 1.5 GB) → ~5–10 min

Usage:
    python -m scripts.scrapers.pbf_processor \
        --pbf osm_data/asia-260327.osm.pbf \
        --continent asia \
        --output scraped_data/pbf_asia.json

    python -m scripts.scrapers.pbf_processor \
        --pbf osm_data/australia-oceania-260327.osm.pbf \
        --continent oceania \
        --output scraped_data/pbf_oceania.json
"""

import argparse
import json
import logging
import math
import os
import sys
import time
from collections import defaultdict
from typing import Optional

import osmium

# Allow running from project root
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from scripts.scrapers.config import TARGET_CITIES

logger = logging.getLogger(__name__)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)

# ---------------------------------------------------------------------------
# Which OSM amenity/tourism/historic tags we care about
# ---------------------------------------------------------------------------

RESTAURANT_AMENITIES = {
    "restaurant", "cafe", "fast_food", "bar", "pub", "food_court",
    "bakery", "ice_cream", "biergarten", "canteen", "food",
}

HOTEL_TOURISM = {"hotel", "hostel", "motel", "guest_house", "apartment", "resort"}
HOTEL_AMENITY = {"hotel", "hostel", "motel"}

ATTRACTION_TOURISM = {
    "museum", "attraction", "viewpoint", "theme_park", "aquarium",
    "zoo", "gallery", "artwork", "camp_site", "caravan_site",
    "picnic_site", "information",
}
ATTRACTION_HISTORIC = {
    "monument", "castle", "ruins", "archaeological_site", "memorial",
    "fort", "palace", "battlefield", "wayside_cross", "wayside_shrine",
    "building", "city_gate", "manor",
}
CULTURAL_AMENITY = {
    "arts_centre", "theatre", "cinema", "concert_hall", "casino",
    "nightclub", "library", "community_centre", "cultural_centre",
    "music_venue",
}

# Leisure tags worth capturing
LEISURE_TAGS = {
    "park", "nature_reserve", "beach", "marina", "water_park",
    "sports_centre", "stadium", "garden", "bird_hide", "hot_spring",
}

# ---------------------------------------------------------------------------
# City bounding boxes (lat_min, lon_min, lat_max, lon_max)
# Generated from TARGET_CITIES coordinates + ±0.5° buffer
# ---------------------------------------------------------------------------

CITY_BBOX_RADIUS = 0.35  # degrees (~35-40 km radius)


def _build_city_index(radius: float = CITY_BBOX_RADIUS) -> list:
    """Return list of (city_config, lat_min, lon_min, lat_max, lon_max)."""
    boxes = []
    for city in TARGET_CITIES:
        lat, lon = city["lat"], city["lon"]
        boxes.append((
            city,
            lat - radius, lon - radius,
            lat + radius, lon + radius,
        ))
    return boxes


def _find_city(lat: float, lon: float, city_index: list) -> Optional[dict]:
    """Return the nearest city config if (lat, lon) falls within any city bbox."""
    for city, lat_min, lon_min, lat_max, lon_max in city_index:
        if lat_min <= lat <= lat_max and lon_min <= lon <= lon_max:
            return city
    return None


# ---------------------------------------------------------------------------
# Tag extractors
# ---------------------------------------------------------------------------

def _tag(tags, *keys, default="") -> str:
    for k in keys:
        v = tags.get(k)
        if v:
            return str(v)
    return default


def _price_from_tags(tags) -> Optional[int]:
    """Map OSM price/stars tags → 1-4 integer price level."""
    p = _tag(tags, "price", "pricerate", "fee").lower()
    mapping = {
        "cheap": 1, "budget": 1, "inexpensive": 1, "low": 1,
        "moderate": 2, "medium": 2, "normal": 2,
        "expensive": 3, "high": 3,
        "luxury": 4, "fine": 4, "very_expensive": 4,
    }
    if p in mapping:
        return mapping[p]
    # Fallback: stars → price level
    stars_raw = _tag(tags, "stars", "hotel:stars")
    try:
        stars = int(float(stars_raw))
        return min(max(math.ceil(stars / 1.25), 1), 4)
    except (ValueError, TypeError):
        pass
    return None


def _dietary_from_tags(tags) -> list:
    """Extract dietary tags from OSM diet:* tags."""
    diets = []
    diet_map = {
        "diet:vegetarian": "vegetarian",
        "diet:vegan": "vegan",
        "diet:halal": "halal",
        "diet:kosher": "kosher",
        "diet:gluten_free": "gluten_free",
    }
    for tag_key, label in diet_map.items():
        val = tags.get(tag_key, "").lower()
        if val in ("yes", "only", "dedicated"):
            diets.append(label)
    return diets


def _cuisine_from_tags(tags) -> list:
    raw = _tag(tags, "cuisine")
    if not raw:
        return []
    return [c.strip() for c in raw.replace(";", ",").split(",") if c.strip()]


def _opening_hours_from_tags(tags) -> str:
    oh = _tag(tags, "opening_hours")
    if not oh or len(oh) > 100:
        return "08:00-22:00"
    return oh


def _extract_restaurant(node_or_way, tags, city: dict) -> dict:
    return {
        "source": "osm_pbf",
        "osm_id": node_or_way,
        "name": _tag(tags, "name", "name:en"),
        "name_en": _tag(tags, "name:en", "name:latin", "name"),
        "name_local": _tag(tags, "name"),
        "destination_id": city["destination_id"],
        "city": city["city"],
        "country": city["country"],
        "country_code": city["country_code"],
        "exchange_rate": city.get("exchange_rate", 25000),
        "lat": None,   # filled after
        "lon": None,
        "cuisine": _cuisine_from_tags(tags),
        "amenity": _tag(tags, "amenity"),
        "opening_hours": _opening_hours_from_tags(tags),
        "phone": _tag(tags, "phone", "contact:phone"),
        "website": _tag(tags, "website", "contact:website"),
        "price_level": _price_from_tags(tags),
        "dietary_tags": _dietary_from_tags(tags),
        "outdoor_seating": tags.get("outdoor_seating", "") == "yes",
        "takeaway": tags.get("takeaway", "") in ("yes", "only"),
        "delivery": tags.get("delivery", "") == "yes",
        "description": _tag(tags, "description", "description:en"),
        "wheelchair": tags.get("wheelchair", ""),
        "addr_street": _tag(tags, "addr:street"),
        "addr_city": _tag(tags, "addr:city"),
    }


def _extract_hotel(osm_id, tags, city: dict) -> dict:
    return {
        "source": "osm_pbf",
        "osm_id": osm_id,
        "name": _tag(tags, "name", "name:en"),
        "name_en": _tag(tags, "name:en", "name"),
        "destination_id": city["destination_id"],
        "city": city["city"],
        "country": city["country"],
        "country_code": city["country_code"],
        "exchange_rate": city.get("exchange_rate", 25000),
        "lat": None,
        "lon": None,
        "property_type": _tag(tags, "tourism", "amenity").upper() or "HOTEL",
        "star_rating": _safe_int(_tag(tags, "stars", "hotel:stars")),
        "rating": None,
        "price_per_night": None,
        "phone": _tag(tags, "phone", "contact:phone"),
        "website": _tag(tags, "website", "contact:website"),
        "email": _tag(tags, "email", "contact:email"),
        "address": _tag(tags, "addr:full", "addr:street"),
        "amenities": _amenities_from_tags(tags),
        "description": _tag(tags, "description", "description:en"),
        "wheelchair": tags.get("wheelchair", ""),
        "internet_access": tags.get("internet_access", ""),
        "rooms": _safe_int(tags.get("rooms")),
    }


def _extract_attraction(osm_id, tags, city: dict, attraction_type: str = "attraction") -> dict:
    return {
        "source": "osm_pbf",
        "osm_id": osm_id,
        "name": _tag(tags, "name", "name:en"),
        "name_en": _tag(tags, "name:en", "name"),
        "name_local": _tag(tags, "name"),
        "destination_id": city["destination_id"],
        "city": city["city"],
        "country": city["country"],
        "country_code": city["country_code"],
        "lat": None,
        "lon": None,
        "type": attraction_type,
        "tourism": _tag(tags, "tourism"),
        "historic": _tag(tags, "historic"),
        "leisure": _tag(tags, "leisure"),
        "amenity": _tag(tags, "amenity"),
        "description": _tag(tags, "description", "description:en"),
        "website": _tag(tags, "website", "contact:website"),
        "opening_hours": _opening_hours_from_tags(tags),
        "fee": tags.get("fee", ""),
        "wheelchair": tags.get("wheelchair", ""),
        "artwork_type": tags.get("artwork_type", ""),
        "museum_type": tags.get("museum:type", ""),
        "heritage": tags.get("heritage", ""),
        "wikipedia": _tag(tags, "wikipedia", "wikidata"),
    }


def _amenities_from_tags(tags) -> list:
    amenities = []
    checks = {
        "wifi": ["internet_access", "wifi"],
        "pool": ["leisure"],
        "parking": ["amenity"],
        "restaurant": ["amenity"],
        "bar": ["amenity"],
    }
    if tags.get("internet_access", "").lower() in ("wlan", "wifi", "yes"):
        amenities.append("Free WiFi")
    if tags.get("swimming_pool") == "yes" or tags.get("leisure") == "swimming_pool":
        amenities.append("Swimming Pool")
    if tags.get("amenity") == "parking" or tags.get("parking") == "yes":
        amenities.append("Parking")
    if tags.get("breakfast") == "yes":
        amenities.append("Breakfast Included")
    if tags.get("air_conditioning") == "yes":
        amenities.append("Air Conditioning")
    return amenities


def _safe_int(val) -> Optional[int]:
    try:
        return int(float(str(val)))
    except (TypeError, ValueError):
        return None


# ---------------------------------------------------------------------------
# osmium Handler
# ---------------------------------------------------------------------------

class WanderlustHandler(osmium.SimpleHandler):
    """
    Streams through OSM PBF and collects nodes/ways matching our categories.
    Filters by city bounding boxes to keep memory usage manageable.
    """

    def __init__(self, city_index: list):
        super().__init__()
        self.city_index = city_index

        # Collected data
        self.restaurants: list = []
        self.hotels: list = []
        self.attractions: list = []

        # Stats
        self._nodes_seen = 0
        self._last_log = time.time()
        self._log_interval = 5_000_000  # log every 5M nodes

    def _log_progress(self, n_type: str):
        self._nodes_seen += 1
        if self._nodes_seen % self._log_interval == 0:
            elapsed = time.time() - self._last_log
            logger.info(
                f"  Processed {self._nodes_seen:,} {n_type}s | "
                f"restaurants={len(self.restaurants):,} "
                f"hotels={len(self.hotels):,} "
                f"attractions={len(self.attractions):,}"
            )
            self._last_log = time.time()

    def node(self, n):
        self._log_progress("node")
        if not n.tags or not n.location.valid():
            return
        lat, lon = n.location.lat, n.location.lon
        city = _find_city(lat, lon, self.city_index)
        if city is None:
            return
        self._process_tags(n.id, n.tags, lat, lon, city, element_type="node")

    def way(self, w):
        self._log_progress("way")
        # Ways (polygons for large POIs like big restaurants, hotels)
        try:
            if not w.tags:
                return
            # Use centroid approximation (first node location if available)
            # pyosmium ways don't expose coords directly — skip for now
            # We get most POIs from nodes
        except Exception:
            pass

    def _process_tags(self, osm_id, tags, lat, lon, city, element_type="node"):
        amenity = tags.get("amenity", "")
        tourism = tags.get("tourism", "")
        historic = tags.get("historic", "")
        leisure = tags.get("leisure", "")
        name = tags.get("name") or tags.get("name:en")

        if not name:
            return  # Skip unnamed POIs

        # --- Restaurants / food ---
        if amenity in RESTAURANT_AMENITIES:
            rec = _extract_restaurant(osm_id, tags, city)
            rec["lat"], rec["lon"] = lat, lon
            self.restaurants.append(rec)

        # --- Hotels ---
        elif tourism in HOTEL_TOURISM or amenity in HOTEL_AMENITY:
            rec = _extract_hotel(osm_id, tags, city)
            rec["lat"], rec["lon"] = lat, lon
            self.hotels.append(rec)

        # --- Tourist attractions ---
        elif tourism in ATTRACTION_TOURISM:
            rec = _extract_attraction(osm_id, tags, city, tourism)
            rec["lat"], rec["lon"] = lat, lon
            self.attractions.append(rec)

        # --- Historic sites ---
        elif historic in ATTRACTION_HISTORIC:
            rec = _extract_attraction(osm_id, tags, city, f"historic_{historic}")
            rec["lat"], rec["lon"] = lat, lon
            self.attractions.append(rec)

        # --- Cultural venues ---
        elif amenity in CULTURAL_AMENITY:
            rec = _extract_attraction(osm_id, tags, city, f"cultural_{amenity}")
            rec["lat"], rec["lon"] = lat, lon
            self.attractions.append(rec)

        # --- Leisure / nature ---
        elif leisure in LEISURE_TAGS:
            rec = _extract_attraction(osm_id, tags, city, leisure)
            rec["lat"], rec["lon"] = lat, lon
            self.attractions.append(rec)


# ---------------------------------------------------------------------------
# Top-level processing function
# ---------------------------------------------------------------------------

def process_pbf(pbf_path: str, output_path: str, continent: str = "unknown"):
    """
    Process a single .osm.pbf file and write JSON output.

    Output JSON schema:
    {
      "meta": { "source_file": ..., "continent": ..., "processed_at": ... },
      "restaurants": [...],
      "hotels": [...],
      "attractions": [...]
    }
    """
    logger.info(f"=== Processing PBF: {pbf_path} ===")
    logger.info(f"File size: {os.path.getsize(pbf_path) / 1e9:.2f} GB")
    logger.info(f"Target cities: {len(TARGET_CITIES)}")
    logger.info("Building city bounding box index...")

    city_index = _build_city_index()
    logger.info(f"  City bboxes: {len(city_index)}")

    handler = WanderlustHandler(city_index)

    logger.info("Starting PBF stream (this may take a while for large files)...")
    start = time.time()

    handler.apply_file(pbf_path, locations=False)

    elapsed = time.time() - start
    logger.info(f"PBF streaming done in {elapsed:.0f}s ({elapsed/60:.1f} min)")
    logger.info(f"  Restaurants : {len(handler.restaurants):,}")
    logger.info(f"  Hotels      : {len(handler.hotels):,}")
    logger.info(f"  Attractions : {len(handler.attractions):,}")

    # Write output
    os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
    output = {
        "meta": {
            "source_file": os.path.basename(pbf_path),
            "continent": continent,
            "processed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "elapsed_seconds": int(elapsed),
            "target_cities": len(TARGET_CITIES),
        },
        "restaurants": handler.restaurants,
        "hotels": handler.hotels,
        "attractions": handler.attractions,
    }

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(output, f, ensure_ascii=False, indent=2)

    size_mb = os.path.getsize(output_path) / 1e6
    logger.info(f"Written: {output_path} ({size_mb:.1f} MB)")

    # Per-city summary
    city_counts: dict = defaultdict(lambda: {"restaurants": 0, "hotels": 0, "attractions": 0})
    for r in handler.restaurants:
        city_counts[r["city"]]["restaurants"] += 1
    for h in handler.hotels:
        city_counts[h["city"]]["hotels"] += 1
    for a in handler.attractions:
        city_counts[a["city"]]["attractions"] += 1

    logger.info("\nPer-city summary:")
    for city_name, counts in sorted(city_counts.items()):
        logger.info(
            f"  {city_name:25s}  "
            f"restaurants={counts['restaurants']:3d}  "
            f"hotels={counts['hotels']:3d}  "
            f"attractions={counts['attractions']:3d}"
        )

    return output


def merge_pbf_output_into_pipeline(pbf_json_path: str):
    """
    Take a processed pbf_*.json file and merge it into the pipeline databases.
    Equivalent to running run_pipeline.py --dry-run=False but from a pre-processed file.
    """
    from scripts.processors import normalize_restaurants, normalize_hotels, normalize_destinations

    with open(pbf_json_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    logger.info(f"Merging {pbf_json_path} into knowledge base...")

    # Restaurants
    raw_restaurants = data.get("restaurants", [])
    if raw_restaurants:
        r_stats = normalize_restaurants.merge_into_cuisine_db(raw_restaurants)
        logger.info(f"Restaurants: +{r_stats['added']} added, {r_stats['skipped_duplicate']} dupes → {r_stats['total_in_db']} total")

    # Attractions → destinations enrichment
    # Group by destination_id for normalize_destinations
    city_groups: dict = defaultdict(lambda: {"destination_id": "", "city": "", "attractions": [], "restaurants": [], "events": []})
    for a in data.get("attractions", []):
        did = a.get("destination_id", "")
        city_groups[did]["destination_id"] = did
        city_groups[did]["city"] = a.get("city", "")
        city_groups[did]["attractions"].append(a)

    if city_groups:
        d_stats = normalize_destinations.enrich_destinations(list(city_groups.values()))
        logger.info(f"Destinations: {d_stats['destinations_enriched']} enriched, +{d_stats['activities_added']} activities")

    # Hotels
    raw_hotels = data.get("hotels", [])
    if raw_hotels:
        # Group by city for normalize_hotels
        hotel_cities: dict = defaultdict(lambda: {"hotels": [], "city": "", "destination_id": ""})
        for h in raw_hotels:
            did = h.get("destination_id", "")
            hotel_cities[did]["hotels"].append(h)
            hotel_cities[did]["city"] = h.get("city", "")
            hotel_cities[did]["destination_id"] = did

        h_stats = normalize_hotels.save_hotels(list(hotel_cities.values()))
        logger.info(f"Hotels: +{h_stats['added']} added, {h_stats['skipped']} skipped → {h_stats['total_hotels']} total")

    logger.info("Merge complete.")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process OSM PBF files for Wanderlust")
    parser.add_argument("--pbf", required=True, help="Path to .osm.pbf file")
    parser.add_argument("--output", required=True, help="Output JSON file path")
    parser.add_argument("--continent", default="unknown", help="Continent name (for metadata)")
    parser.add_argument("--merge", action="store_true",
                        help="After processing, merge into knowledge base databases")
    args = parser.parse_args()

    result = process_pbf(args.pbf, args.output, args.continent)

    if args.merge:
        merge_pbf_output_into_pipeline(args.output)

    print(f"\nDone! {len(result['restaurants'])} restaurants, {len(result['hotels'])} hotels, {len(result['attractions'])} attractions")
    print(f"Output: {args.output}")
    if args.merge:
        print("Merged into knowledge base.")
    else:
        print(f"To merge: python -m scripts.scrapers.pbf_processor --pbf {args.pbf} --output {args.output} --merge")