import contextlib
import fcntl
import hashlib
import json
import os
from pathlib import Path
import resource
import subprocess
import time

ROOT = Path(__file__).resolve().parents[1]
RUNS = ROOT / "runs"
VENDOR = ROOT / "vendor"
LIMIT = 10 * 1024**3

def read(path):
    return json.loads(Path(path).read_text())

def write(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n")
    tmp.replace(path)

def sha(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for b in iter(lambda: f.read(1024 * 1024), b""):
            h.update(b)
    return h.hexdigest()

def fingerprint(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True).encode()).hexdigest()

def immutable(path, settings):
    if Path(path).exists() and read(path) != settings:
        raise RuntimeError(f"Settings changed: {path}. Use a new run directory.")
    write(path, settings)

@contextlib.contextmanager
def exclusive():
    RUNS.mkdir(exist_ok=True)
    with (RUNS / ".training.lock").open("w") as f:
        try:
            fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise RuntimeError("Another kit command is running. Run one job at a time.")
        yield

def memory():
    # macOS ru_maxrss is bytes; current RSS includes mapped CPU memory.
    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    rss = int(subprocess.check_output(["ps", "-o", "rss=", "-p", str(os.getpid())])) * 1024
    return {"rss_bytes": rss, "peak_rss_bytes": peak}

def guard(extra=0):
    m = memory()
    # GPU allocations may overlap RSS. max avoids counting unified allocations twice.
    if max(m["rss_bytes"], extra) > LIMIT:
        raise MemoryError("10 GiB memory target exceeded; checkpoint retained. Reduce the batch size.")
    return m

def jsonl(path, rows):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    with tmp.open("w") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n")
    tmp.replace(path)

def rows(path):
    return [json.loads(x) for x in Path(path).read_text().splitlines() if x.strip()]

def assert_disjoint(splits, keys):
    owner = {}
    for split, records in splits.items():
        for row in records:
            for key in keys:
                value = row.get(key)
                if value is None:
                    continue
                token = (key, json.dumps(value, sort_keys=True))
                if token in owner and owner[token] != split:
                    raise ValueError(f"Cross-split overlap: {key} in {owner[token]} and {split}")
                owner[token] = split

def paired_ci(before, after, metric, ratio=False, seed=44, repeats=2000):
    """Paired cluster bootstrap; ratio aggregates edits/reference units, not clip WER."""
    import numpy as np
    a, b = {r["id"]: r for r in before}, {r["id"]: r for r in after}
    if a.keys() != b.keys() or not a:
        raise ValueError("Paired evaluation must have identical, nonempty IDs")
    groups = {}
    for key, row in a.items():
        group = row.get("group", key)
        if group != b[key].get("group", key):
            raise ValueError("Paired cluster mismatch")
        groups.setdefault(group, []).append(key)
    av, bv, den = [], [], []
    for ids in groups.values():
        av.append(sum(a[k][metric] for k in ids))
        bv.append(sum(b[k][metric] for k in ids))
        if ratio and any(a[k]["units"] != b[k]["units"] for k in ids):
            raise ValueError("Reference units changed")
        den.append(sum(a[k]["units"] for k in ids) if ratio else len(ids))
    av, bv, den = map(np.asarray, (av, bv, den))
    rng = np.random.default_rng(seed)
    changes = []
    for _ in range(repeats):
        ix = rng.integers(0, len(av), len(av))
        changes.append(float((bv[ix].sum() - av[ix].sum()) / den[ix].sum()))
    return {"delta_candidate_minus_baseline": float((bv.sum()-av.sum())/den.sum()),
            "ci95": np.quantile(changes, [.025, .975]).tolist(), "clusters": len(groups),
            "bootstrap_repeats": repeats}

def fleet_score(base, candidate):
    if set(base) != set(candidate) or "ko" not in base or len(base) != 8:
        raise ValueError("Expected all eight languages")
    eligible = candidate["ko"]["error_rate"] <= base["ko"]["error_rate"]
    eligible &= all(candidate[l]["error_rate"] <= base[l]["error_rate"] + .005 for l in base if l != "ko")
    changes = [(base[l]["error_rate"] - candidate[l]["error_rate"]) / max(base[l]["error_rate"], 1e-12) for l in base]
    return bool(eligible), sum(changes) / len(changes)

def stable_object_splits(groups, old, confirmation=100):
    assignment, new = {}, []
    for group, records in groups.items():
        prior = {old[r["metadata"]["asset_id"]] for r in records if r["metadata"]["asset_id"] in old}
        if len(prior)>1:
            raise ValueError("New duplicate links bridge original splits; cannot preserve split isolation")
        if prior:
            assignment[group] = prior.pop()
        else:
            new.append(group)
    new.sort(key=lambda s: hashlib.sha256(("44/"+s).encode()).hexdigest())
    if len(new)<confirmation*2:
        raise ValueError("Insufficient disjoint new groups for expansion and confirmation")
    for i, group in enumerate(new):
        assignment[group] = "confirmation" if i<confirmation else ("validation" if (i-confirmation)%5==0 else "train")
    return assignment