# -*- coding: utf-8 -*-
"""试验设计正式 tc 套件：TestClient 走真实 API + 落库 block_no 结构断言（增广/α-格子）。

fixture（镜像 e2e_trial_design_ext）：16 种质池 G01..G16 + mk_study helper
（TrialStudyModel block_count + TrialStudyEntryModel entry_number/germplasm_id）。
经真实 HTTP 端点 POST /api/v1/bre/statistics/trial-design 断言：
  [1] augmented：12 条目（9 新品系 + 3 对照）× 3 区组 → 新品系每区组 3 条不重复、对照每区组
      重复（check_reps=3）且对照 block_no=None（不落库）、新品系 block_no∈{1,2,3}
  [2] alpha 平方 v=9=k²（k=3 素数）：r=4 完全格子 → 4 rep×3 块×3 条、每 rep 覆盖全部 9、
      任一对至多同块一次（λ≤1）、balance_warning=False、block_no=rep*100+block 落库
  [3] alpha 部分 r=3（同 v=9）：仍可分解平衡、λ≤1
  [4] alpha 平方 v=16=k²（k=4 非素数）：r=5 → balance_warning=True（可分解）
  [5] alpha 一般 v=12（非平方，k=3 显式）：s=4 贪心可分解、r 缺省=s=4、block_no 落库
  [6] 校验 409：对照非条目子集 / 无对照 / v 不被 k 整除 / 非平方缺 block_size
DB 读回统一放在 TestClient 退出后的单一 asyncio.run（避免嵌套事件循环）。
依赖: Redis + PG 正常（TestClient 走真实 lifespan）。运行后自动清理。
"""
import os
os.environ["ENVIRONMENT"] = "dev"
os.environ["PYTHONUTF8"] = "1"

import sys, asyncio  # noqa: E402
sys.path.insert(0, r"d:\dpb\dpb\backend")

import main  # noqa: E402
from fastapi.testclient import TestClient  # noqa: E402
from sqlalchemy import delete, select  # noqa: E402

from app.core.database import create_async_engine_and_session  # noqa: E402
from app.api.v1.module_system.user.model import UserModel  # noqa: E402 (注册 mapper)
from app.api.v1.module_bre.target.model import TargetModel  # noqa: E402
from app.api.v1.module_bre.germplasm.model import BreedingGermplasmModel  # noqa: E402
from app.api.v1.module_bre.cross_combination.model import CrossCombinationModel  # noqa: E402
from app.api.v1.module_bre.trial.model import TrialModel  # noqa: E402
from app.api.v1.module_bre.trial_study.model import TrialStudyModel  # noqa: E402
from app.api.v1.module_bre.trial_study_entry.model import TrialStudyEntryModel  # noqa: E402

create_app = main.create_app
TOKEN = None
ok, fail = 0, 0

PREFIX = "TCTDES"
tokens: dict[str, list[int]] = {
    "target": [], "germ": [], "combo": [], "trial": [], "study": [], "entry": [],
}
FIX: dict = {}
RES: dict = {}


def check(name, cond, detail=""):
    global ok, fail
    if cond:
        ok += 1
        print(f"  [ok] {name} {detail}")
    else:
        fail += 1
        print(f"  [FAIL] {name} {detail}")


def pairs_of(blk_entries):
    ps = []
    for i in range(len(blk_entries)):
        for j in range(i + 1, len(blk_entries)):
            ps.append(tuple(sorted((blk_entries[i], blk_entries[j]))))
    return ps


def login(client):
    global TOKEN
    d = {"username": "super", "password": "123456", "grant_type": "password", "login_type": "PC端"}
    r = client.post("/api/v1/system/auth/login", data=d)
    b = r.json()
    if r.status_code == 200 and b.get("code") == 0:
        TOKEN = b["data"]["access_token"]
        return
    key = client.get("/api/v1/system/auth/captcha/get").json()["data"]["key"]
    client.post("/api/v1/system/auth/captcha/slider/complete", json={"captcha_key": key})
    d["captcha_key"] = key
    r = client.post("/api/v1/system/auth/login", data=d)
    b = r.json()
    assert r.status_code == 200 and b.get("code") == 0, f"LOGIN FAIL {r.status_code} {b}"
    TOKEN = b["data"]["access_token"]


def auth():
    return {"Authorization": f"Bearer {TOKEN}"}


async def _wipe() -> None:
    """启动前清理本前缀残留（防上次进程被杀/建 fixture 中途失败留下的脏数据）。"""
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            study_ids = list((await db.execute(select(TrialStudyModel.id).where(
                TrialStudyModel.study_name.like(f"{PREFIX}-%")))).scalars())
            if study_ids:
                entries = list((await db.execute(select(TrialStudyEntryModel.id).where(
                    TrialStudyEntryModel.trial_study_id.in_(study_ids)))).scalars())
                if entries:
                    await db.execute(delete(TrialStudyEntryModel).where(
                        TrialStudyEntryModel.id.in_(entries)))
                await db.execute(delete(TrialStudyModel).where(TrialStudyModel.id.in_(study_ids)))
            trial_ids = list((await db.execute(select(TrialModel.id).where(
                TrialModel.trial_name == f"试验{PREFIX}"))).scalars())
            if trial_ids:
                await db.execute(delete(TrialModel).where(TrialModel.id.in_(trial_ids)))
            combo_ids = list((await db.execute(select(CrossCombinationModel.id).where(
                CrossCombinationModel.combination_code == f"{PREFIX}-C"))).scalars())
            if combo_ids:
                await db.execute(delete(CrossCombinationModel).where(
                    CrossCombinationModel.id.in_(combo_ids)))
            germ_ids = list((await db.execute(select(BreedingGermplasmModel.id).where(
                BreedingGermplasmModel.cultivar_name.like(f"{PREFIX}-G%")))).scalars())
            if germ_ids:
                await db.execute(delete(BreedingGermplasmModel).where(
                    BreedingGermplasmModel.id.in_(germ_ids)))
            await db.execute(delete(TargetModel).where(TargetModel.target_name == f"{PREFIX}-T"))
            await db.commit()
            print(f"[preclean] trial_design tc 前缀残留已清（study={len(study_ids)} germ={len(germ_ids)}）")
    finally:
        await engine.dispose()


async def _build_fixture() -> None:
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            target = TargetModel(target_name=f"{PREFIX}-T", created_id=1)
            db.add(target)
            await db.flush()
            tokens["target"].append(target.id)

            combo = CrossCombinationModel(
                combination_code=f"{PREFIX}-C", bre_target_id=target.id,
                female_parent_id=None, male_parent_id=None,
                design_type="full_diallel", created_id=1)
            db.add(combo)
            await db.flush()
            tokens["combo"].append(combo.id)

            germs = []
            for j in range(1, 17):
                g = BreedingGermplasmModel(cultivar_name=f"{PREFIX}-G{j:02d}",
                                           can_be_female=(j % 2 == 0), can_be_male=True, created_id=1)
                db.add(g)
                germs.append(g)
            await db.flush()
            tokens["germ"] = [g.id for g in germs]
            gid = [g.id for g in germs]

            trial = TrialModel(trial_name=f"试验{PREFIX}", design_type="rcbd", created_id=1)
            db.add(trial)
            await db.flush()
            tokens["trial"].append(trial.id)

            async def mk_study(name, block_count, n_entries, germ_subset):
                st = TrialStudyModel(trial_id=trial.id, study_name=f"{PREFIX}-{name}",
                                     block_count=block_count, design_type="rcbd", year=2025, created_id=1)
                db.add(st)
                await db.flush()
                ents = []
                for e in range(1, n_entries + 1):
                    ent = TrialStudyEntryModel(trial_study_id=st.id, entry_number=e,
                                               germplasm_id=germ_subset[e - 1], created_id=1)
                    db.add(ent)
                    ents.append(ent)
                await db.flush()
                return st, ents

            s1, ents1 = await mk_study("AUG", 3, 12, gid[0:12])
            s3, ents3 = await mk_study("ALPHA9", 3, 9, gid[0:9])
            s2, ents2 = await mk_study("ALPHA16", 4, 16, gid[0:16])
            s4, ents4 = await mk_study("ALPHA12", 4, 12, gid[0:12])
            tokens["study"] += [s1.id, s3.id, s2.id, s4.id]
            tokens["entry"] += [e.id for e in ents1 + ents3 + ents2 + ents4]
            FIX["studies"] = {"AUG": s1.id, "ALPHA9": s3.id, "ALPHA16": s2.id, "ALPHA12": s4.id}
            FIX["gid"] = gid
            await db.commit()
    finally:
        await engine.dispose()


async def _block_map(db, sid: int) -> dict:
    rows = (await db.execute(
        select(TrialStudyEntryModel.entry_number, TrialStudyEntryModel.block_no)
        .where(TrialStudyEntryModel.trial_study_id == sid))).all()
    return {r.entry_number: r.block_no for r in rows}


async def _readback_and_cleanup() -> None:
    """TestClient 退出后：读回 block_no 断言落库结构，然后清理。"""
    engine, sf = create_async_engine_and_session()
    try:
        async with sf() as db:
            # ===== AUG 读回：对照不落库、新品系 block_no∈{1,2,3} =====
            aug_bn = await _block_map(db, FIX["studies"]["AUG"])
            check("[1] 对照(entry 1..3) block_no=None 不落库",
                  all(aug_bn.get(e) is None for e in (1, 2, 3)),
                  f"{ {e: aug_bn.get(e) for e in (1, 2, 3)} }")
            check("[1] 新品系(entry 4..12) block_no∈{1,2,3}",
                  all(aug_bn.get(e) in (1, 2, 3) for e in range(4, 13)),
                  f"{ {e: aug_bn.get(e) for e in (4, 13)} }")

            # ===== ALPHA9 / ALPHA12 读回：复合编码 rep*100+block =====
            a9_bn = await _block_map(db, FIX["studies"]["ALPHA9"])
            exp3 = {e: int(r) * 100 + int(b) for r, bm in RES["a9"]["blocks"].items()
                    for b, vs in bm.items() for e in vs}
            check("[2] α9 block_no=rep*100+block 落库",
                  all(a9_bn.get(e) == exp3[e] for e in exp3) and len(a9_bn) == 9,
                  f"{dict(list(a9_bn.items())[:4])}…")

            a12_bn = await _block_map(db, FIX["studies"]["ALPHA12"])
            exp4 = {e: int(r) * 100 + int(b) for r, bm in RES["a12"]["blocks"].items()
                    for b, vs in bm.items() for e in vs}
            check("[5] α12(一般 v) block_no=rep*100+block 落库",
                  all(a12_bn.get(e) == exp4[e] for e in exp4) and len(a12_bn) == 12)

            # ===== 清理 =====
            if tokens["entry"]:
                await db.execute(delete(TrialStudyEntryModel).where(
                    TrialStudyEntryModel.id.in_(tokens["entry"])))
            if tokens["study"]:
                await db.execute(delete(TrialStudyModel).where(TrialStudyModel.id.in_(tokens["study"])))
            if tokens["trial"]:
                await db.execute(delete(TrialModel).where(TrialModel.id.in_(tokens["trial"])))
            if tokens["combo"]:
                await db.execute(delete(CrossCombinationModel).where(
                    CrossCombinationModel.id.in_(tokens["combo"])))
            if tokens["germ"]:
                await db.execute(delete(BreedingGermplasmModel).where(
                    BreedingGermplasmModel.id.in_(tokens["germ"])))
            if tokens["target"]:
                await db.execute(delete(TargetModel).where(TargetModel.id.in_(tokens["target"])))
            await db.commit()
            print(f"[cleanup] trial_design tc 数据已清（entry={len(tokens['entry'])} study={len(tokens['study'])}）")
    finally:
        await engine.dispose()


def main_() -> None:
    asyncio.run(_wipe())
    asyncio.run(_build_fixture())
    try:
        with TestClient(create_app()) as client:
            login(client)
            H = auth()
            sid = FIX["studies"]
            gid = FIX["gid"]

            def run(study_id, design_type, **kw):
                payload = {"trial_study_id": study_id, "design_type": design_type, **kw}
                r = client.post("/api/v1/bre/statistics/trial-design", json=payload, headers=H)
                check(f"[HTTP] trial-design {design_type} 200", r.status_code == 200,
                      f"{r.status_code} {str(r.text)[:120]}")
                b = r.json()
                return b.get("data") if b.get("code") == 0 else None

            # ================= [1] 增广设计 =================
            checks = gid[0:3]
            aug = run(sid["AUG"], "augmented", seed=7, check_germplasm_ids=checks)
            RES["aug"] = aug
            check("[1] 输出结构 n_new=9/n_checks=3/check_reps=3/block=3",
                  aug is not None and aug["design_type"] == "augmented" and aug["n_new_entries"] == 9
                  and aug["n_checks"] == 3 and aug["check_reps"] == 3
                  and aug["block_count"] == 3 and aug["n_entries"] == 12,
                  f"{aug and (aug['n_new_entries'], aug['n_checks'], aug['check_reps'])}")
            if aug:
                check("[1] 新品系均分 3 区组×3 条、恰为 {4..12} 不重复",
                      set(aug["blocks"].keys()) == {"1", "2", "3"}
                      and all(len(v) == 3 for v in aug["blocks"].values())
                      and set(e for v in aug["blocks"].values() for e in v) == set(range(4, 13)),
                      f"{aug['blocks']}")
                check("[1] checks 恰 3 对照 {entry_number,germplasm_id}",
                      len(aug["checks"]) == 3 and {c["germplasm_id"] for c in aug["checks"]} == set(checks)
                      and all(set(c.keys()) == {"entry_number", "germplasm_id"} for c in aug["checks"]),
                      f"{aug['checks']}")

            # ================= [3] α-格子部分 r=3（同 v=9，先跑：r=4 会整表覆写 block_no） =================
            full9 = list(range(1, 10))
            a9p = run(sid["ALPHA9"], "alpha", seed=11, reps=3)
            if a9p:
                b9p = a9p["blocks"]
                cnt3 = {}
                for bm in b9p.values():
                    for v in bm.values():
                        for p in pairs_of(v):
                            cnt3[p] = cnt3.get(p, 0) + 1
                check("[3] 部分 r=3：3 rep×3 块×3 条、可分解、仍平衡 λ≤1",
                      a9p["reps"] == 3 and set(b9p.keys()) == {"1", "2", "3"}
                      and all(sorted(e for v in bm.values() for e in v) == full9 for bm in b9p.values())
                      and all(c <= 1 for c in cnt3.values()),
                      f"reps={a9p['reps']} maxλ={max(cnt3.values()) if cnt3 else 0}")

            # ================= [2] α-格子平方 v=9 k=3（素数，完全格子 r=4，λ≤1，最后跑确保落库=全量） =================
            a9 = run(sid["ALPHA9"], "alpha", seed=11)
            RES["a9"] = a9
            if a9:
                b9 = a9["blocks"]
                check("[2] v=9 自动 k=3/r=4/s=3/warn=False",
                      a9["block_size"] == 3 and a9["reps"] == 4 and a9["block_count"] == 3
                      and a9["balance_warning"] is False and a9["n_entries"] == 9,
                      f"k={a9['block_size']} r={a9['reps']} s={a9['block_count']} warn={a9['balance_warning']}")
                check("[2] 4 rep×3 块×3 条",
                      set(b9.keys()) == {"1", "2", "3", "4"}
                      and all(set(bm.keys()) == {"1", "2", "3"} for bm in b9.values())
                      and all(len(v) == 3 for bm in b9.values() for v in bm.values()))
                check("[2] 每 rep 可分解（覆盖全部 9）",
                      all(sorted(e for v in bm.values() for e in v) == full9 for bm in b9.values()))
                cnt = {}
                for bm in b9.values():
                    for v in bm.values():
                        for p in pairs_of(v):
                            cnt[p] = cnt.get(p, 0) + 1
                check("[2] 任一对至多同块一次（λ≤1）", all(c <= 1 for c in cnt.values()) and len(cnt) == 36,
                      f"maxλ={max(cnt.values()) if cnt else 0}")

            # ================= [4] α-格子平方 v=16 k=4（非素数 → 平衡保守告警） =================
            a16 = run(sid["ALPHA16"], "alpha", seed=13)
            full16 = list(range(1, 17))
            if a16:
                b16 = a16["blocks"]
                check("[4] v=16 自动 k=4/r=5/s=4/warn=True（k 非素数）",
                      a16["block_size"] == 4 and a16["reps"] == 5 and a16["block_count"] == 4
                      and a16["balance_warning"] is True,
                      f"k={a16['block_size']} r={a16['reps']} warn={a16['balance_warning']}")
                check("[4] 5 rep×4 块×4 条 + 每 rep 可分解",
                      set(b16.keys()) == {"1", "2", "3", "4", "5"}
                      and all(set(bm.keys()) == {"1", "2", "3", "4"} for bm in b16.values())
                      and all(len(v) == 4 for bm in b16.values() for v in bm.values())
                      and all(sorted(e for v in bm.values() for e in v) == full16 for bm in b16.values()))

            # ================= [5] α-格子一般 v=12（非平方，k=3 显式，贪心可分解） =================
            a12 = run(sid["ALPHA12"], "alpha", seed=17, block_size=3)
            RES["a12"] = a12
            full12 = list(range(1, 13))
            if a12:
                b12 = a12["blocks"]
                check("[5] v=12 显式 k=3/s=4/r 缺省=4",
                      a12["block_size"] == 3 and a12["block_count"] == 4 and a12["reps"] == 4
                      and a12["n_entries"] == 12,
                      f"k={a12['block_size']} s={a12['block_count']} r={a12['reps']}")
                check("[5] 4 rep×4 块×3 条 + 每 rep 可分解（贪心未破坏分解性）",
                      set(b12.keys()) == {"1", "2", "3", "4"}
                      and all(set(bm.keys()) == {"1", "2", "3", "4"} for bm in b12.values())
                      and all(len(v) == 3 for bm in b12.values() for v in bm.values())
                      and all(sorted(e for v in bm.values() for e in v) == full12 for bm in b12.values()),
                      f"warn={a12['balance_warning']}")

            # ================= [6] 校验 409 =================
            r = client.post("/api/v1/bre/statistics/trial-design",
                            json={"trial_study_id": sid["AUG"], "design_type": "augmented",
                                  "seed": 1, "check_germplasm_ids": [gid[14]]}, headers=H)
            check("[6] augmented 对照非条目子集 → 409", r.status_code == 409, f"{r.status_code}")

            r = client.post("/api/v1/bre/statistics/trial-design",
                            json={"trial_study_id": sid["AUG"], "design_type": "augmented",
                                  "seed": 1, "check_germplasm_ids": []}, headers=H)
            check("[6] augmented 无对照 → 409", r.status_code == 409, f"{r.status_code}")

            r = client.post("/api/v1/bre/statistics/trial-design",
                            json={"trial_study_id": sid["ALPHA16"], "design_type": "alpha",
                                  "seed": 1, "block_size": 5}, headers=H)
            check("[6] alpha v 不被 k 整除 → 409", r.status_code == 409, f"{r.status_code}")

            r = client.post("/api/v1/bre/statistics/trial-design",
                            json={"trial_study_id": sid["ALPHA12"], "design_type": "alpha", "seed": 1},
                            headers=H)
            check("[6] alpha 非平方缺 block_size → 409", r.status_code == 409, f"{r.status_code}")
    finally:
        asyncio.run(_readback_and_cleanup())

    print(f"\n===== trial_design tc 套件：ok={ok} fail={fail} =====")


if __name__ == "__main__":
    main_()
    sys.exit(1 if fail else 0)