"""
外部进程多实例全局分摊 —— distribute_quota 逻辑单测（无需网关/缓存后端）。

用 dict 作缓存桩，模拟同类型多个外部进程实例注册，验证扩缩容时全局 count 按
「存活实例数 N」分摊（而非每个进程都全额），以及实例退出/租约过期后的接管行为。
"""
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../utils/worker/src"))

import chongming_worker.runtime as r

PASS = 0
FAIL = 0


def report(name, cond, detail=""):
    global PASS, FAIL
    if cond:
        PASS += 1
        print(f"✅ {name}" + (f" — {detail}" if detail else ""))
    else:
        FAIL += 1
        print(f"❌ {name}" + (f" — {detail}" if detail else ""))


def fake_cache():
    return {}


def as_instance(iid, cache, worker_type="demo"):
    """把当前身份设为 iid，并在缓存登记该实例参与方。"""
    old = os.environ.get("CHONGMING_WORKER_INSTANCE")
    os.environ["CHONGMING_WORKER_INSTANCE"] = iid
    try:
        r.register_instance(worker_type, cache)
    finally:
        if old is None:
            os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
        else:
            os.environ["CHONGMING_WORKER_INSTANCE"] = old


def quota_as(iid, worker_type, total, cache):
    """以指定实例身份计算其分摊配额（external 分支）。"""
    old = os.environ.get("CHONGMING_WORKER_INSTANCE")
    os.environ["CHONGMING_WORKER_INSTANCE"] = iid
    try:
        return r.distribute_quota(cache, worker_type, total, gw_id=None)
    finally:
        if old is None:
            os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
        else:
            os.environ["CHONGMING_WORKER_INSTANCE"] = old


def unreg_instance(iid, cache, worker_type="demo"):
    """以指定实例身份从参与方列表移除它。"""
    old = os.environ.get("CHONGMING_WORKER_INSTANCE")
    os.environ["CHONGMING_WORKER_INSTANCE"] = iid
    try:
        r.unregister_instance(worker_type, cache)
    finally:
        if old is None:
            os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
        else:
            os.environ["CHONGMING_WORKER_INSTANCE"] = old


def run():
    cache = fake_cache()
    wt = "demo"

    # 1. 单实例（未启用实例注册，参与方为空）→ 兜底全额（兼容旧行为）
    q = r.distribute_quota(cache, wt, 6, gw_id=None)
    report("单实例无参与方 → 全额承担 6", q == 6, f"got={q}")

    # 2. 两个实例并存：全局 6 分摊 → 各 3
    as_instance("inst-a", cache, wt)
    as_instance("inst-b", cache, wt)
    qa = quota_as("inst-a", wt, 6, cache)
    qb = quota_as("inst-b", wt, 6, cache)
    report("两个实例并存 → 各分 3", qa == 3 and qb == 3, f"a={qa}, b={qb}")

    # 3. 三个实例、总数不能整除：7 // 3 余 1 → 按序给第一个补 1
    as_instance("inst-c", cache, wt)
    qs = {i: quota_as(i, wt, 7, cache) for i in ("inst-a", "inst-b", "inst-c")}
    report(
        "三个实例 ≠整除 → 总和=7 且余量补靠前实例",
        sum(qs.values()) == 7 and set(qs.values()) == {2, 2, 3},
        f"{qs}",
    )

    # 4. 一个实例退出（unregister）→ 其余接管全部
    unreg_instance("inst-c", cache, wt)
    qa2 = quota_as("inst-a", wt, 6, cache)
    qb2 = quota_as("inst-b", wt, 6, cache)
    report(
        "实例退出后由剩余接管 → 各分 3",
        qa2 == 3 and qb2 == 3,
        f"a={qa2}, b={qb2}",
    )

    # 5. 租约过期清理：伪造一个过期的参与方，读实例列表应排除它
    cache[r._SCALE_INSTANCES_KEY.format(type=wt)] = [
        ["inst-stale", time.time() - r._INSTANCE_LEASE - 10],
        ["inst-fresh", time.time()],
    ]
    ids = r._read_instances(cache, wt)
    report(
        "死实例按租约剔除",
        ids == ["inst-fresh"],
        f"instances={ids}",
    )

    # 6. 实例数=1 时全额（不误分摊）
    cache[r._SCALE_INSTANCES_KEY.format(type=wt)] = [
        [r.current_instance_id(), time.time()]
    ]
    q1 = r.distribute_quota(cache, wt, 5, gw_id=None)
    report("单存活实例 → 全额 5", q1 == 5, f"got={q1}")

    # 7. inline（有网关身份）路径回归：按在线网关分摊不受影响
    cache[r._SCALE_GATEWAYS_KEY] = [["gw-a", time.time()], ["gw-b", time.time()]]
    qa_in = r.distribute_quota(cache, wt, 6, "gw-a")
    qb_in = r.distribute_quota(cache, wt, 6, "gw-b")
    report(
        "inline 按网关分摊回归 → 各 3",
        qa_in == 3 and qb_in == 3,
        f"gw-a={qa_in}, gw-b={qb_in}",
    )

    # 8. 逐实例绝对覆盖：写入后能被读回且供启动解析（external）使用
    r._write_own_instance_count(cache, wt, 4, "inst-a")
    r._write_own_instance_count(cache, wt, 2, "inst-b")
    report(
        "instcounts 覆盖写入/读回",
        r._own_instance_count(cache, wt, "inst-a") == 4
        and r._own_instance_count(cache, wt, "inst-b") == 2,
        f"{r._read_instance_counts(cache, wt)}",
    )
    # 启动解析：external（gw_id=None）有绝对覆盖时直接采用，而非被全局均分
    old = os.environ.get("CHONGMING_WORKER_INSTANCE")
    os.environ["CHONGMING_WORKER_INSTANCE"] = "inst-a"
    try:
        restored = r.resolve_worker_count(wt, 6, cache, "cache", gw_id=None)
    finally:
        if old is None:
            os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
        else:
            os.environ["CHONGMING_WORKER_INSTANCE"] = old
    report("resolve_worker_count 启动时采用本实例绝对数", restored == 4, f"got={restored}")

    # 9. 清除覆盖后回退全局均分
    r._clear_instance_counts(cache, wt)
    report(
        "清除覆盖后覆盖表为空",
        r._read_instance_counts(cache, wt) == {},
        f"{r._read_instance_counts(cache, wt)}",
    )

    # 10. 缺省实例身份唯一（hostname:pid）：同机未指定 --instance 的多进程不会并成一个参与方
    r.set_current_instance_id(None)
    old_env = os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
    try:
        default_iid = r.current_instance_id()
    finally:
        if old_env is not None:
            os.environ["CHONGMING_WORKER_INSTANCE"] = old_env
        else:
            os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
    import socket as _socket
    report(
        "缺省实例身份含进程区分（{host}:{pid}）",
        default_iid == f"{_socket.gethostname()}:{os.getpid()}",
        f"got={default_iid}",
    )

    print("\n" + "=" * 56)
    print(f"测试结果: ✅ {PASS} 通过  ❌ {FAIL} 失败")
    return 0 if FAIL == 0 else 1


if __name__ == "__main__":
    sys.exit(run())