"""Phase 4.3 — shared-user-task personalization overlap sweep.

Stress-tests the personalization claim (per-user memory vs. a live global/shared
store) as users' task sets *overlap*. Workload: every user has ``tasks_per_user``
signature tasks with a FIXED per-user argument realization (user A always NYC,
user B always Boston), consistent across the user's sessions. A fraction
``overlap_frac`` of those tasks come from ONE shared template pool that every
user reuses (same templates, different per-user args); the rest are private and
disjoint. As overlap rises, a global store sees many users' different calls for
the same template — does per-user memory still win, or does a shared store catch
up?

Two arms, both live + evicting (only partitioning differs), so the contrast is
personalization alone:
  - global_evict  : one live global store (no per-user view)
  - personal_memory (ours): per-user live evicting store

Targets are the served model's greedy calls (real, generated once per unique
query on the migrated node and cached per overlap level). Run from ``code/``:
    python -m harness.phase4_overlap --url http://localhost:30000/v1
"""
from __future__ import annotations

import argparse
import json
import os
from collections import defaultdict
from pathlib import Path

from . import metrics
from .client import ToolClient
from .data import load_bfcl
from .memory import Embedder, GlobalEvict, PersonalMemory
from .run_accept import _parse_target, generate_targets
from .simulate import build_users

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"
# Tokenizer for the token-LCP accept metric: HF hub id by default;
# override with a local snapshot path if running offline.
MODEL_PATH = os.environ.get("SPECMEM_TOKENIZER", "openai/gpt-oss-120b")


class Slot:
    """generate_targets expects objects with .query and .functions; the
    functions registry is identical across BFCL tasks' first-tool schema, so we
    attach it lazily below."""
    __slots__ = ("query", "functions")

    def __init__(self, query, functions=None):
        self.query = query
        self.functions = functions


def _replay(instances, targets, embedder, capacity):
    arms = [GlobalEvict(capacity=40 * capacity),
            PersonalMemory(capacity=capacity, eviction="lru")]
    agg = {a.name: defaultdict(list) for a in arms}
    for ins in instances:
        tgt = targets.get(ins.query)
        if tgt is None:
            continue
        for a in arms:
            agg[a.name][ins.session].append(metrics.score(
                a.draft(ins.query, ins.functions, ins.user_id, embedder), tgt))
        cn, ca = _parse_target(tgt)
        for a in arms:
            a.observe(ins.query, ins.functions, ins.user_id, cn, ca, embedder)
    out = {}
    for n, v in agg.items():
        post = [x for s, xs in v.items() if s > 0 for x in xs]
        out[n] = round(sum(x["accept_length"] for x in post) /
                       max(1, len(post)), 3)
    return out


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--url", default="http://localhost:30000/v1")
    p.add_argument("--model", default="gpt-oss-120b")
    p.add_argument("--workers", type=int, default=16)
    p.add_argument("--overlaps", nargs="+", type=float,
                   default=[0.0, 0.25, 0.5, 0.75, 1.0])
    p.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2])
    p.add_argument("--capacity", type=int, default=48)
    args = p.parse_args()

    metrics.get_tokenizer(MODEL_PATH)
    tasks = load_bfcl()
    embedder = Embedder()
    client = ToolClient(url=args.url, model=args.model)
    if not client.ping():
        raise SystemExit(f"served model not reachable at {args.url}")

    sweep = {}
    for ov in args.overlaps:
        per_seed = {"global_evict": [], "personal_memory": []}
        for sd in args.seeds:
            instances = build_users(
                tasks, n_users=40, tasks_per_user=15, n_sessions=12,
                queries_per_session=6, seed=sd, overlap_frac=ov,
                user_consistent=True)
            # attach the (shared) functions registry to Slots for target gen
            fmap = {ins.query: ins.functions for ins in instances}
            cache_f = RESULTS / f"phase4_overlap_targets_ov{ov}_seed{sd}.json"
            cache = json.loads(cache_f.read_text()) if cache_f.exists() else {}
            miss = [Slot(q, fmap[q]) for q in {i.query for i in instances}
                    if q not in cache]
            if miss:
                cache.update(generate_targets(client, miss, workers=args.workers))
                cache_f.write_text(json.dumps(cache))
            res = _replay(instances, cache, embedder, args.capacity)
            per_seed["global_evict"].append(res["global_evict"])
            per_seed["personal_memory"].append(res["personal_memory"])
            print(f"ov={ov} seed={sd}: {res}", flush=True)
        import statistics as st
        g = st.mean(per_seed["global_evict"])
        pm = st.mean(per_seed["personal_memory"])
        sweep[ov] = {
            "global_evict": round(g, 3),
            "personal_memory": round(pm, 3),
            "personal_advantage_pct": round(100 * (pm - g) / g, 1),
            "per_seed": per_seed}
        print(f"== ov={ov}: global {g:.3f} | personal {pm:.3f} | "
              f"adv {sweep[ov]['personal_advantage_pct']}% ==", flush=True)

    out = {"config": {"users": 40, "tasks_per_user": 15, "sessions": 12,
                      "capacity": args.capacity, "seeds": args.seeds,
                      "user_consistent_args": True,
                      "arms": "global_evict (live shared) vs personal_memory (ours)"},
           "sweep": sweep}
    (RESULTS / "phase4_personalization_overlap.json").write_text(
        json.dumps(out, indent=2))
    print(json.dumps({ov: {"global": s["global_evict"],
                           "personal": s["personal_memory"],
                           "adv%": s["personal_advantage_pct"]}
                      for ov, s in sweep.items()}, indent=1))


if __name__ == "__main__":
    main()