import os
import time
import requests
import psycopg2

GEOCODE = "https://geocode.search.hereapi.com/v1/geocode"
MATRIX = "https://matrix.router.hereapi.com/v8"
API_KEY = os.environ["HERE_API_KEY"]

K = 10   # shortlist size. Bounds the matrix, and therefore the bill.


# --- Stage 1: geocode the customer. Cached permanently. -------------------
_geo_cache: dict[str, dict] = {}


def geocode_customer(address: str) -> dict:
    """Buildings are stationary. Key on the NORMALIZED address."""
    key = normalize(address)
    if key in _geo_cache:
        return _geo_cache[key]

    resp = requests.get(GEOCODE, params={
        "q": address, "limit": 1, "apiKey": API_KEY,
    }, timeout=10)
    if resp.status_code == 403:
        raise RuntimeError("Geocoding not entitled")
    resp.raise_for_status()

    items = resp.json().get("items", [])
    if not items:
        raise ValueError("Address not found")

    it = items[0]
    # `access` is the road-network entry point. Route to this, not `position`.
    pos = (it.get("access") or [it["position"]])[0]
    _geo_cache[key] = pos
    return pos


def normalize(addr: str) -> str:
    return " ".join(addr.lower().replace(",", " ").split())


# --- Stage 2: PostGIS shortlist. ZERO API calls. --------------------------
def shortlist(conn, lat: float, lng: float, k: int = K) -> list[dict]:
    """
    k nearest stores by straight-line distance.

    Sub-millisecond against a GIST index. Indifferent to your traffic.
    This step exists to make stage 3 cheap.
    """
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, name, ST_Y(geom::geometry), ST_X(geom::geometry)
            FROM stores
            WHERE open_now
            ORDER BY geom <-> ST_SetSRID(ST_MakePoint(%s, %s), 4326)::geography
            LIMIT %s
        """, (lng, lat, k))          # note: MakePoint is (lng, lat)
        return [
            {"id": r[0], "name": r[1], "lat": r[2], "lng": r[3]}
            for r in cur.fetchall()
        ]


# --- Stage 3: ONE matrix call. Not k routing calls. -----------------------
def rank_by_drive_time(origin: dict, stores: list[dict]) -> list[dict]:
    """
    1 origin × k destinations = ONE matrix request.

    k routing calls would be k requests. At store-locator traffic,
    that ratio is your entire cost structure.
    """
    body = {
        "origins": [{"lat": origin["lat"], "lng": origin["lng"]}],
        "destinations": [{"lat": s["lat"], "lng": s["lng"]} for s in stores],
        "regionDefinition": {
            # Region mode. Max 400km diameter. 1×10 is well within sync limits.
            "type": "circle",
            "center": {"lat": origin["lat"], "lng": origin["lng"]},
            "radius": 100_000,
        },
        "matrixAttributes": ["travelTimes", "distances"],
    }

    resp = requests.post(f"{MATRIX}/matrix",
                         params={"apiKey": API_KEY, "async": "false"},
                         json=body, timeout=30)
    if resp.status_code == 403:
        raise RuntimeError("Matrix Routing not entitled")
    resp.raise_for_status()

    m = resp.json()["matrix"]
    n_dest = m["numDestinations"]

    ranked = []
    for j, store in enumerate(stores):
        # Flat, row-major. origin index is 0. Write the helper once.
        idx = 0 * n_dest + j
        err = (m.get("errorCodes") or [0] * n_dest)[idx]
        if err:
            continue    # 1=disconnected 2=no match 3=violates restriction
        ranked.append({
            **store,
            "travel_time_s": m["travelTimes"][idx],
            "distance_m": m["distances"][idx],
        })

    return sorted(ranked, key=lambda s: s["travel_time_s"])


# --- Put together ---------------------------------------------------------
def nearest_stores(conn, address: str) -> list[dict]:
    customer = geocode_customer(address)
    candidates = shortlist(conn, customer["lat"], customer["lng"])
    if not candidates:
        return []                       # Do NOT silently expand the radius
    return rank_by_drive_time(customer, candidates)


if __name__ == "__main__":
    conn = psycopg2.connect(os.environ["DATABASE_URL"])
    for s in nearest_stores(conn, "425 W Randolph St, Chicago, IL"):
        mins = round(s["travel_time_s"] / 60)
        print(f"{mins} min — {s['name']}")