#!/usr/bin/env python3
"""
Shared helpers for reading MSMARCO-XI. One place, so every script agrees.

THE ACTUAL PASSAGES LAYOUT (confirmed from the real file)
--------------------------------------------------------
    passages: Struct({
        'English_passages'   : List(String),
        'Translated_passages': List(String),
        'is_selected'        : List(Int64),
    })

This is a STRUCT OF PARALLEL LISTS, not a list of structs. One row holds one
dict whose three values are equal-length lists, aligned by position:

    row["English_passages"][i]  <->  row["Translated_passages"][i]
                                <->  row["is_selected"][i]

That index `i` is exactly the `passage_index` half of the canonical ID, so the
same passage keeps the same identity in every language file. There is NO url
field, so pseudo-document grouping falls back to query_id.

`iter_passages` normalises both layouts so callers never branch on it.
"""
from __future__ import annotations

import os
from pathlib import Path
from typing import Iterator

# MSMARCO-XI language codes are ISO-639-3 + script, e.g. "asm_Beng", "npi_Deva".
ISO3_TO_ISO2 = {
    "asm": "as", "ben": "bn", "guj": "gu", "hin": "hi", "kan": "kn",
    "mal": "ml", "mar": "mr",
    "nep": "ne", "npi": "ne",          # dataset uses npi_Deva for Nepali
    "ory": "or", "ori": "or",
    "pan": "pa", "san": "sa", "tam": "ta", "tel": "te", "urd": "ur",
    "eng": "en",
}

LANG_NAMES = {
    "as": "Assamese", "bn": "Bengali", "gu": "Gujarati", "hi": "Hindi",
    "kn": "Kannada", "ml": "Malayalam", "mr": "Marathi", "ne": "Nepali",
    "or": "Odia", "pa": "Punjabi", "sa": "Sanskrit", "ta": "Tamil",
    "te": "Telugu", "ur": "Urdu", "en": "English",
}


def norm_lang(raw) -> str:
    """'asm_Beng' -> 'as'.  Unknown 2-letter codes pass through."""
    if not raw:
        return "en"
    code = str(raw).split("_")[0].lower()
    return ISO3_TO_ISO2.get(code, code if len(code) == 2 else code)


def iter_passages(row, t_key: str, en_key: str | None,
                  sel_key: str | None, url_key: str | None
                  ) -> Iterator[tuple[int, str | None, str | None, int, str | None]]:
    """Yield (index, text, text_en, is_selected, url) for either layout.

    Layout A - struct of parallel lists (MSMARCO-XI):
        {"Translated_passages": [...], "English_passages": [...], "is_selected": [...]}
    Layout B - list of structs (the common alternative):
        [{"passage_text": ..., "is_selected": ...}, ...]
    """
    def at(seq, i):
        return seq[i] if isinstance(seq, (list, tuple)) and i < len(seq) else None

    if isinstance(row, dict):
        texts = row.get(t_key)
        if isinstance(texts, (list, tuple)):                       # layout A
            ens = row.get(en_key) if en_key else None
            sels = row.get(sel_key) if sel_key else None
            urls = row.get(url_key) if url_key else None
            for i, txt in enumerate(texts):
                yield (i, txt, at(ens, i), int(at(sels, i) or 0), at(urls, i))
        else:                                                      # single struct
            yield (0, texts,
                   row.get(en_key) if en_key else None,
                   int(row.get(sel_key) or 0) if sel_key else 0,
                   row.get(url_key) if url_key else None)

    elif isinstance(row, (list, tuple)):                           # layout B
        for i, item in enumerate(row):
            if isinstance(item, dict):
                yield (i, item.get(t_key),
                       item.get(en_key) if en_key else None,
                       int(item.get(sel_key) or 0) if sel_key else 0,
                       item.get(url_key) if url_key else None)


def detect_layout(row) -> str:
    if isinstance(row, dict):
        if any(isinstance(v, (list, tuple)) for v in row.values()):
            return "struct_of_lists"
        return "single_struct"
    if isinstance(row, (list, tuple)) and row and isinstance(row[0], dict):
        return "list_of_structs"
    return "unknown"


def default_root() -> Path:
    """Data root, resolved without relying on the environment.

    Order: $VOICERAG_ROOT  ->  <repo>/../voicerag_data

    The second form is derived from this file's location, so it is correct on
    jupyter-pod (/workspace/carbine/anurag) and kls-headnode
    (/home/dgx-i-carbine/carbine/anurag) alike, and in a fresh shell where
    nothing has been exported.
    """
    env = os.environ.get("VOICERAG_ROOT")
    if env:
        return Path(env).expanduser().resolve()
    return (Path(__file__).resolve().parents[2] / "voicerag_data").resolve()


def resolve_files(report: dict, root: Path) -> tuple[list[Path], list[str]]:
    """Map file paths recorded in schema_report.json onto the CURRENT data root.

    The same NFS export is mounted at different points on the two hosts:
        jupyter-pod   /workspace/carbine/anurag/voicerag_data
        kls-headnode  /home/dgx-i-carbine/carbine/anurag/voicerag_data

    A report written on one host therefore holds paths that do not exist on the
    other. Rather than force a re-run of schema inspection, we rebase: try the
    recorded path, then root-relative, then splice the path from its `hf_cache`
    segment onto the current root.

    Returns (existing_paths, unresolved_strings).
    """
    found: list[Path] = []
    missing: list[str] = []
    for f in report.get("files", []):
        p = Path(f)
        if p.exists():
            found.append(p)
            continue
        cand = root / f                       # relative form (written by newer runs)
        if cand.exists():
            found.append(cand)
            continue
        parts = p.parts                       # absolute form from the other host
        if "hf_cache" in parts:
            cand = root.joinpath(*parts[parts.index("hf_cache"):])
            if cand.exists():
                found.append(cand)
                continue
        missing.append(f)
    return found, missing


def load_report(root: Path) -> dict:
    """Read schema_report.json and rebase its file list onto `root`."""
    import json
    rp = root / "data" / "schema_report.json"
    if not rp.exists():
        raise SystemExit(
            f"\nMissing prerequisite: {rp}\n\n"
            "Run first:  python scripts/02_inspect_schema.py --root " + str(root) + "\n"
        )
    rep = json.loads(rp.read_text())
    files, missing = resolve_files(rep, root)
    if not files:
        raise SystemExit(
            f"\nNone of the {len(rep.get('files', []))} parquet paths in {rp.name} exist "
            f"under {root}.\n"
            f"  first recorded path: {rep.get('files', ['<none>'])[0]}\n\n"
            "The report was probably written on the other host. Re-run:\n"
            f"  python scripts/02_inspect_schema.py --root {root}\n"
        )
    if missing:
        print(f"  note: {len(missing)} of {len(rep['files'])} recorded paths "
              f"could not be resolved (using {len(files)})")
    rep["files"] = [str(f) for f in files]
    return rep

def pick_device(explicit: str | None = None) -> str:
    """Which torch device to use. $VOICERAG_DEVICE wins.

    It has to win, because on a ZeroGPU Space autodetection is actively wrong.
    Free-tier Gradio Spaces are ZeroGPU-only, and OUTSIDE a @spaces.GPU function
    ZeroGPU enables a CUDA *emulation* so models can be placed on 'cuda' during
    startup. torch.cuda.is_available() therefore returns True while no real GPU
    is attached, and autodetect would put bge-m3 on a device that cannot run a
    forward pass.

    app.py declares one @spaces.GPU function -- ZeroGPU refuses to start a Space
    without one -- and never calls it. Quota is billed per call, so the Space
    consumes none of the free tier's 5-minute daily allowance and can serve
    unlimited requests, instead of going dark when a judge exhausts it.
    """
    import os

    if explicit:
        return explicit
    env = os.environ.get("VOICERAG_DEVICE", "").strip().lower()
    if env:
        return env
    import torch

    return "cuda" if torch.cuda.is_available() else "cpu"