import os
import psycopg2

COUNTRY_FALLBACK = {
    "FR": "France",
    "DE": "Germany",
    "IT": "Italy",
    "GB": "United Kingdom",
    "ES": "Spain",
    "SE": "Sweden",
    "NO": "Norway",
    "FI": "Finland",
    "PL": "Poland",
    "NL": "Netherlands",
    "BE": "Belgium",
    "CH": "Switzerland",
    "AT": "Austria",
    "PT": "Portugal",
    "GR": "Greece",
    "CZ": "Czech Republic",
    "SK": "Slovakia",
    "SI": "Slovenia",
    "HU": "Hungary",
    "UA": "Ukraine",
    "RU": "Russia",
    "DK": "Denmark",
}

def make_query(location_text: str, country_code: str) -> str:
    loc = (location_text or "").strip()
    cc = (country_code or "").strip().upper()
    ctry = COUNTRY_FALLBACK.get(cc, cc)
    return f"{loc}, {ctry}" if loc else ctry

def main():
    dsn = os.environ["DATABASE_URL"]
    batch = int(os.environ.get("GEOCODE_ENQUEUE_BATCH", "5000"))

    # Retry policy
    retry_after_hours = int(os.environ.get("GEOCODE_RETRY_AFTER_HOURS", "72"))  # 72h = 3 jours
    force_retry = os.environ.get("GEOCODE_FORCE_RETRY", "0").strip() in ("1", "true", "yes", "on")

    conn = psycopg2.connect(dsn)
    conn.autocommit = True

    with conn.cursor() as cur:
        # On sélectionne :
        # - bands non géocodés avec location_text
        # - soit pas encore en queue, soit queue en error/queued et next_run_at <= now()
        # - et si erreur récente: on respecte retry_after_hours sauf si force_retry
        cur.execute(
            """
            SELECT
              b.ma_id, b.country, b.location_text
            FROM bands b
            LEFT JOIN geocode_queue q ON q.ma_id = b.ma_id
            WHERE b.geom IS NULL
              AND b.location_text IS NOT NULL AND b.location_text <> ''
              AND (
                   q.ma_id IS NULL
                   OR (q.status IN ('queued','error') AND q.next_run_at <= now())
              )
              AND (
                   %s
                   OR b.geocode_error IS NULL
                   OR b.geocode_error_at IS NULL
                   OR b.geocode_error_at < now() - (%s || ' hours')::interval
              )
            ORDER BY
              COALESCE(q.next_run_at, now()) ASC,
              b.ma_id ASC
            LIMIT %s
            """,
            (force_retry, retry_after_hours, batch),
        )

        rows = cur.fetchall()
        enq = 0

        for ma_id, country, location_text in rows:
            qtxt = make_query(location_text, country)

            cur.execute(
                """
                INSERT INTO geocode_queue (ma_id, query, country, status, next_run_at)
                VALUES (%s, %s, %s, 'queued', now())
                ON CONFLICT (ma_id) DO UPDATE
                  SET query = EXCLUDED.query,
                      country = EXCLUDED.country,
                      status = CASE
                        WHEN geocode_queue.status IN ('done','processing') THEN geocode_queue.status
                        ELSE 'queued'
                      END,
                      next_run_at = CASE
                        WHEN geocode_queue.status IN ('done','processing') THEN geocode_queue.next_run_at
                        ELSE now()
                      END,
                      updated_at = now()
                """,
                (ma_id, qtxt, country),
            )
            enq += 1

        print(f"[enqueue] queued_or_updated={enq} (selected={len(rows)}) force_retry={force_retry} retry_after_hours={retry_after_hours}")

    conn.close()

if __name__ == "__main__":
    main()