"""Unified, resumable embedding stage — replaces model-2's embed_train.py + embed_test.py.

    python src/embed.py --split train --part pool --gpu 0
    python src/embed.py --split test  --part pool --gpu 1
    python src/embed.py --split test  --part s1   --gpu 1

Parts
-----
  pool : Source 2 + Source 3 (S2 rows first, then S3) -- the search corpus, doc prefix
  s1   : Source 1 -- the queries, query prefix
  all  : pool then s1

Notes
-----
* `--split train --part s1` is OPTIONAL and off by default in scripts/queue.sh.
  run_train_eval.py only queries n_s1_ce+n_s1_stack+n_s1_val (=200k) of the 2.21M
  train S1, and caches those itself. Embedding all 2.21M is the dead work that cost
  model-2 ~50 min. `--split test --part s1` IS required (every test S1 must be scored).
* Every run is resumable and fingerprint-checked -- see encoder.py.
* --gpu is mandatory-by-default so nothing silently lands on a busy cuda:0.
"""
from __future__ import annotations

import argparse
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import pandas as pd                                     # noqa: E402
from config import CFG, EMB, PATHS, emb_dir, model_spec  # noqa: E402
from encoder import encode_to_shards, write_ids          # noqa: E402

COLS = ["entity_id", "business_name", "business_address", "country"]
T0 = time.time()


def log(m):
    print(f"[{time.time() - T0:7.0f}s] {m}", flush=True)


def record_text(name, addr, country):
    """Must stay byte-identical to the text used at search time."""
    return f"{name} | {addr} | {country}"


def read_part(split, part):
    p = PATHS[split]
    if part == "pool":
        df = pd.concat([pd.read_parquet(p["source2"], columns=COLS),
                        pd.read_parquet(p["source3"], columns=COLS)], ignore_index=True)
    else:
        df = pd.read_parquet(p["source1"], columns=COLS)
    for c in COLS:
        df[c] = df[c].fillna("").astype(str)
    return df


def run(split, part, device, model_name=None):
    hf_id, qpre, dpre, _dim, lic = model_spec(model_name)
    prefix = dpre if part == "pool" else qpre
    df = read_part(split, part)
    log(f"{split}/{part}: {len(df):,} records | model={hf_id} ({lic}) | prefix={prefix!r}")
    texts = [record_text(a, b, c) for a, b, c in
             zip(df["business_name"], df["business_address"], df["country"])]
    ids = df["entity_id"].tolist()
    del df
    d = emb_dir(f"{split}_{part}", model_name)
    os.makedirs(d, exist_ok=True)
    write_ids(os.path.join(d, "ids.txt"), ids)
    return encode_to_shards(texts, ids, d, hf_id, prefix, EMB["max_len"], EMB["batch"],
                            device, EMB["shard_rows"], log)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--split", choices=["train", "test"], required=True)
    ap.add_argument("--part", choices=["pool", "s1", "all"], default="pool")
    ap.add_argument("--gpu", default=None,
                    help="GPU index, or 'cpu'. Required unless M3_GPU is set -- this is "
                         "deliberate: model-2 lost runs by defaulting to a saturated cuda:0.")
    ap.add_argument("--model", default=None, help=f"override (default {EMB['model']})")
    a = ap.parse_args()

    gpu = a.gpu if a.gpu is not None else os.environ.get("M3_GPU")
    if gpu is None:
        sys.exit("refusing to guess a GPU: pass --gpu 0 / --gpu 1 / --gpu cpu (or set M3_GPU)")
    device = "cpu" if str(gpu) == "cpu" else f"cuda:{int(gpu)}"

    parts = ["pool", "s1"] if a.part == "all" else [a.part]
    for p in parts:
        run(a.split, p, device, a.model)
    log("all requested parts complete")


if __name__ == "__main__":
    main()