"""
Export HAS_FINGERPRINT_RULE relationships to CSV via APOC, using ID-range
pagination, and automatic resume on failure.

Usage:
  # Single process (full range)
  python export_hfr.py

  # Auto-split into N parallel workers (each gets its own ID range + state file)
  python export_hfr.py --workers 4

  # Manual range (useful to re-run a specific slice)
  python export_hfr.py --start 0 --end 500000
"""

import argparse
import json
import math
import os
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

from neo4j import GraphDatabase

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

NEO4J_URI = os.getenv("NEO4J_URI")
NEO4J_USER = os.getenv("NEO4J_USER")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")
NEO4J_DATABASE = "graph.db"

# One of: Vendor | Subaccount | Track
ANCHOR = "Track:Orchard"

# Rows per batch (controls both ID-range window and APOC's internal flush size)
BATCH_SIZE = 100_000

ANCHOR_SLUG = ANCHOR.lower().replace(":", "_")
FILE_PREFIX = f"hfr_{ANCHOR_SLUG}"

# ---------------------------------------------------------------------------
# Inner query — label interpolated at build time (labels can't be parameterized)
# $lastId and $nextLastId are passed via APOC's params config key
# ---------------------------------------------------------------------------

INNER_QUERY = f"""
MATCH (x:{ANCHOR}) WHERE x.id > $lastId AND x.id <= $nextLastId
WITH x ORDER BY x.id
MATCH (x)-[hfr:HAS_FINGERPRINT_RULE]->(fpr:FingerprintRule)
RETURN
  x.id          AS orchard_obj_id,
  fpr.territory AS territory,
  fpr.service   AS service,
  fpr.policy    AS policy,
  hfr.createdAt       AS createdAt,
  hfr.createdBy       AS createdBy,
  hfr.lastModifiedAt  AS lastModifiedAt,
  hfr.lastModifiedBy  AS lastModifiedBy,
  hfr.start     AS period_start
""".strip()

# ---------------------------------------------------------------------------
# State helpers
# ---------------------------------------------------------------------------

def state_file(worker_id: int | None) -> Path:
    suffix = f"_w{worker_id}" if worker_id is not None else ""
    return Path(f"export_state_{ANCHOR_SLUG}{suffix}.json")


def load_state(worker_id: int | None, range_start: int) -> dict:
    sf = state_file(worker_id)
    if sf.exists():
        state = json.loads(sf.read_text())
        label = f"worker {worker_id}" if worker_id is not None else "main"
        print(f"[{label}] Resuming: batch={state['batch']}, last_id={state['last_id']}", flush=True)
        return state
    return {"last_id": range_start, "batch": 0}


def save_state(worker_id: int | None, last_id: int, batch: int, done: bool = False) -> None:
    state_file(worker_id).write_text(json.dumps({"last_id": last_id, "batch": batch, "done": done}))


def is_done(worker_id: int | None) -> bool:
    sf = state_file(worker_id)
    return sf.exists() and json.loads(sf.read_text()).get("done", False)


# ---------------------------------------------------------------------------
# Queries
# ---------------------------------------------------------------------------

def find_batch_boundary(session, last_id: int, range_end: int) -> tuple[int | None, int]:
    result = session.run(
        f"""
        MATCH (x:{ANCHOR}) WHERE EXISTS {{ (x)-[:HAS_FINGERPRINT_RULE]->() }}
          AND x.id > $lastId AND x.id <= $rangeEnd
        WITH x ORDER BY x.id LIMIT $batchSize
        RETURN max(x.id) AS nextLastId, count(x) AS cnt
        """,
        lastId=last_id,
        rangeEnd=range_end,
        batchSize=BATCH_SIZE,
    )
    record = result.single()
    if not record or record["cnt"] == 0:
        return None, 0
    return record["nextLastId"], record["cnt"]


def run_export_batch(
    session, last_id: int, next_last_id: int, batch: int, worker_id: int | None
) -> dict:
    suffix = f"_w{worker_id}" if worker_id is not None else ""
    filename = f"{FILE_PREFIX}{suffix}_batch_{batch:04d}.csv"
    result = session.run(
        """
        CALL apoc.export.csv.query(
          $cypher,
          $filename,
          {
            batchSize: $batchSize,
            stream: false,
            params: {lastId: $lastId, nextLastId: $nextLastId}
          }
        )
        YIELD file, rows, time, done
        RETURN file, rows, time, done
        """,
        cypher=INNER_QUERY,
        filename=filename,
        batchSize=BATCH_SIZE,
        lastId=last_id,
        nextLastId=next_last_id,
    )
    return dict(result.single())


# ---------------------------------------------------------------------------
# Single-worker export loop
# ---------------------------------------------------------------------------

def run_worker(range_start: int, range_end: int, worker_id: int | None = None) -> int:
    label = f"worker {worker_id}" if worker_id is not None else "main"
    state = load_state(worker_id, range_start)
    last_id: int = state["last_id"]
    batch: int = state["batch"]
    total_rows = 0

    driver = GraphDatabase.driver(
        NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD),
        connection_timeout=30, connection_acquisition_timeout=30,
    )
    try:
        with driver.session(database=NEO4J_DATABASE) as session:
            while True:
                next_last_id, anchor_count = find_batch_boundary(session, last_id, range_end)

                if next_last_id is None:
                    print(f"[{label}] Done — {total_rows:,} total rows exported.", flush=True)
                    save_state(worker_id, last_id, batch, done=True)
                    break

                print(
                    f"[{label}] Batch {batch:04d}: {anchor_count:,} anchors "
                    f"id ({last_id}, {next_last_id}]",
                    end=" ... ",
                    flush=True,
                )

                row = run_export_batch(session, last_id, next_last_id, batch, worker_id)
                total_rows += row["rows"]
                print(f"{row['rows']:,} rows, {row['time']}ms -> {row['file']}", flush=True)

                last_id = next_last_id
                batch += 1
                save_state(worker_id, last_id, batch)

    except Exception as exc:
        print(f"\n[{label}] Failed on batch {batch}: {exc}", file=sys.stderr)
        print(f"[{label}] State saved — re-run to resume from id > {last_id}", file=sys.stderr)
        raise
    finally:
        driver.close()

    return total_rows


# ---------------------------------------------------------------------------
# Range pre-calculation
# ---------------------------------------------------------------------------

def compute_ranges(n_workers: int) -> list[tuple[int, int]]:
    driver = GraphDatabase.driver(
        NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD),
        connection_timeout=30, connection_acquisition_timeout=30,
    )
    try:
        with driver.session(database=NEO4J_DATABASE) as session:
            record = session.run(
                f"MATCH (x:{ANCHOR}) RETURN min(x.id) AS minId, max(x.id) AS maxId"
            ).single()
            min_id, max_id = record["minId"], record["maxId"]
    finally:
        driver.close()

    print(f"ID range: [{min_id}, {max_id}] — splitting into {n_workers} workers", flush=True)
    chunk = math.ceil((max_id - min_id + 1) / n_workers)
    ranges = []
    for i in range(n_workers):
        start = min_id - 1 + i * chunk  # exclusive lower bound
        end = min(min_id - 1 + (i + 1) * chunk, max_id)
        ranges.append((start, end))
    return ranges


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=1, help="Number of parallel workers (default: 1)")
    parser.add_argument("--start", type=int, default=None, help="Exclusive lower bound ID for manual range")
    parser.add_argument("--end", type=int, default=None, help="Inclusive upper bound ID for manual range")
    args = parser.parse_args()

    if args.workers > 1 and (args.start is not None or args.end is not None):
        print("--start/--end cannot be combined with --workers > 1", file=sys.stderr)
        sys.exit(1)

    if args.workers > 1:
        ranges = compute_ranges(args.workers)
        for i, (s, e) in enumerate(ranges):
            status = "DONE (skipping)" if is_done(i) else f"id ({s}, {e}]"
            print(f"  worker {i}: {status}")

        pending = [(i, s, e) for i, (s, e) in enumerate(ranges) if not is_done(i)]
        if not pending:
            print("All workers already completed.")
            sys.exit(0)

        failed = False
        with ProcessPoolExecutor(max_workers=len(pending)) as pool:
            futures = {
                pool.submit(run_worker, s, e, i): i
                for i, s, e in pending
            }
            for future in as_completed(futures):
                worker_id = futures[future]
                try:
                    future.result()
                except Exception as exc:
                    print(f"Worker {worker_id} failed: {exc}", file=sys.stderr)
                    failed = True
        if failed:
            sys.exit(1)

    else:
        range_start = args.start if args.start is not None else 0
        range_end = args.end if args.end is not None else 2**62
        try:
            run_worker(range_start, range_end, worker_id=None)
        except Exception:
            sys.exit(1)


if __name__ == "__main__":
    main()