import os
import time
import requests

BASE = "https://matrix.router.hereapi.com/v8"
API_KEY = os.environ["HERE_API_KEY"]


class EntitlementError(Exception): ...


def submit_matrix(origins: list[dict], destinations: list[dict],
                  region_center: dict, radius_m: int = 200_000) -> tuple[str, str]:
    """
    Submit an async matrix job.

    Returns (matrixId, statusUrl).
    PERSIST BOTH BEFORE POLLING. If your worker restarts mid-poll, the job is
    still running on HERE's side. Resubmitting bills again.
    """
    body = {
        "origins": origins,            # [{"lat": .., "lng": ..}, ...]
        "destinations": destinations,
        "regionDefinition": {          # Region mode: max 400km diameter
            "type": "circle",
            "center": region_center,
            "radius": radius_m,
        },
        "matrixAttributes": ["travelTimes", "distances"],
    }

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

    data = resp.json()
    # Use the statusUrl HERE returns. Do NOT construct it — global load
    # balancing means a self-built URL may resolve to a different region.
    return data["matrixId"], data["statusUrl"]


def poll_until_done(status_url: str, max_wait_s: int = 600) -> str:
    """Poll with backoff. Returns the result URL on success."""
    backoff, waited = 2.0, 0.0
    while waited < max_wait_s:
        resp = requests.get(status_url, params={"apiKey": API_KEY},
                            timeout=15, allow_redirects=False)
        if resp.status_code in (301, 302, 303):
            return resp.headers["Location"]

        resp.raise_for_status()
        status = resp.json().get("status")
        if status == "completed":
            return resp.json()["resultUrl"]
        if status == "failed":
            raise RuntimeError(f"Matrix failed: {resp.json()}")

        time.sleep(backoff)
        waited += backoff
        backoff = min(backoff * 1.5, 30.0)

    raise TimeoutError(f"Matrix not complete after {max_wait_s}s")


def fetch_result(result_url: str) -> dict:
    resp = requests.get(result_url, params={"apiKey": API_KEY}, timeout=60)
    resp.raise_for_status()
    return resp.json()


# ---------------------------------------------------------------------------
# THE INDEXING HELPER. Write this once. Unit-test it. Never index inline.
# ---------------------------------------------------------------------------
def cell(matrix: dict, origin_idx: int, dest_idx: int, attr: str = "travelTimes"):
    """
    Result arrays are FLAT and ROW-MAJOR.

    A 2x2 returns travelTimes: [73, 1231, 983, 400]
      [0][0]=73  [0][1]=1231
      [1][0]=983 [1][1]=400

    Index it wrong and you get a silently transposed matrix. Every travel time
    is plausible. Every assignment is wrong. Nothing throws.
    """
    n_dest = matrix["numDestinations"]
    return matrix[attr][origin_idx * n_dest + dest_idx]


def cell_error(matrix: dict, origin_idx: int, dest_idx: int) -> int | None:
    """
    Per-cell error codes. Read them.
      1 = graph disconnected   2 = waypoint matching failed
      3 = route found but VIOLATES A RESTRICTION
      4 = waypoint outside region
    """
    codes = matrix.get("errorCodes")
    if not codes:
        return None
    n_dest = matrix["numDestinations"]
    code = codes[origin_idx * n_dest + dest_idx]
    return code or None


if __name__ == "__main__":
    depots = [{"lat": 41.8845, "lng": -87.6386},
              {"lat": 41.5250, "lng": -88.0817}]
    stops = [{"lat": 41.9028, "lng": -87.6317},
             {"lat": 41.7508, "lng": -87.7713}]

    matrix_id, status_url = submit_matrix(depots, stops,
                                          region_center={"lat": 41.8, "lng": -87.8})
    # PERSIST matrix_id AND status_url HERE, before polling.
    result_url = poll_until_done(status_url)
    m = fetch_result(result_url)["matrix"]

    for i in range(m["numOrigins"]):
        for j in range(m["numDestinations"]):
            err = cell_error(m, i, j)
            if err == 3:
                print(f"depot {i} → stop {j}: route VIOLATES a restriction")
            elif err:
                print(f"depot {i} → stop {j}: error {err}")
            else:
                print(f"depot {i} → stop {j}: {cell(m, i, j)}s")

const BASE = "https://matrix.router.hereapi.com/v8";
const API_KEY = process.env.HERE_API_KEY;

const sleep = (ms) => new Promise((r) => setTimeout(r, ms));

export async function submitMatrix(origins, destinations, regionCenter, radiusM = 200_000) {
  const resp = await fetch(`${BASE}/matrix?apiKey=${API_KEY}`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      origins,
      destinations,
      regionDefinition: { type: "circle", center: regionCenter, radius: radiusM },
      matrixAttributes: ["travelTimes", "distances"],
    }),
  });

  if (resp.status === 403) throw new Error("Matrix Routing not entitled");
  if (!resp.ok) throw new Error(`HERE returned ${resp.status}`);

  const { matrixId, statusUrl } = await resp.json();
  // Persist BOTH before polling.
  return { matrixId, statusUrl };
}

export async function pollUntilDone(statusUrl, maxWaitMs = 600_000) {
  let backoff = 2000, waited = 0;
  while (waited < maxWaitMs) {
    const resp = await fetch(`${statusUrl}?apiKey=${API_KEY}`, { redirect: "manual" });
    if ([301, 302, 303].includes(resp.status)) return resp.headers.get("location");

    const body = await resp.json();
    if (body.status === "completed") return body.resultUrl;
    if (body.status === "failed") throw new Error(JSON.stringify(body));

    await sleep(backoff);
    waited += backoff;
    backoff = Math.min(backoff * 1.5, 30_000);
  }
  throw new Error("Matrix timed out");
}

/** Flat, row-major. Write once, test once, never inline. */
export function cell(matrix, originIdx, destIdx, attr = "travelTimes") {
  return matrix[attr][originIdx * matrix.numDestinations + destIdx];
}