"""大批直通良率专用回刷脚本(gxerp)

只解析检验批的大批归属(BatchSt.zt_batch)并重算各大批的 直通_* 统计,
不重建批次的其他统计数据(那个用 correct_batchst.py)。

用法:
    python scripts/correct_zt_gxerp.py [起始日期] [--fresh] [--workers=N] [--skip-resolve]
    起始日期默认 2026-06-01; --fresh 忽略旧锁重新解析归属并强制重算白料(纠错模式);
    --workers 并发线程数, 默认 8; --skip-resolve 跳过归属解析, 直接用已落库的 zt_batch 算大批
"""
import os
import sys
import django

CUR_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_DIR = os.path.dirname(CUR_DIR)
sys.path.insert(0, BASE_DIR)

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "server.settings")
django.setup()

from django.conf import settings
settings.DEBUG = False  # DEBUG 下 Django 会把每条 SQL 记入内存, 长脚本越跑越慢

import concurrent.futures
import threading
from datetime import datetime
from django.db.models import Q
from apps.wpm.models import BatchSt
from apps.wpm.scripts.batch_gxerp import resolve_zt_big, cal_zt_big
import json
from apps.utils.tools import MyJSONEncoder

_print_lock = threading.Lock()


def main(since: str = "2026-06-01", fresh: bool = False, workers: int = 8, skip_resolve: bool = False):
    since_dt = datetime.strptime(since, "%Y-%m-%d")
    done = [0]

    if skip_resolve:
        # 直接用已落库的归属
        seen = set(BatchSt.objects.filter(
            create_time__gte=since_dt, version=1, zt_batch__isnull=False,
        ).values_list("zt_batch", flat=True).distinct())
        bigs = sorted(seen)
        print(f"跳过归属解析, 已落库大批 {len(bigs)} 个 (workers={workers})", flush=True)
    else:
        cond = Q(data__has_key="尺寸检验_日期") | Q(data__has_key="外观检验_日期")
        if fresh:
            # 纠错模式需覆盖曾被(错误)锁定/计算过的批
            cond = cond | Q(zt_batch__isnull=False) | Q(data__has_key="直通_良率")
        qs = BatchSt.objects.filter(create_time__gte=since_dt, version=1).filter(cond).order_by("create_time")
        items = list(qs)
        total = len(items)
        print(f"检验批 {total} 个, 开始解析大批归属 (fresh={fresh}, workers={workers})", flush=True)

        def resolve_one(bs):
            try:
                # fresh 模式忽略旧锁, 从谱系重新解析(纠正多级拆批错锁)
                big = resolve_zt_big(bs, bs.data or {}, use_lock=not fresh)
                if big and bs.zt_batch != big:
                    bs.zt_batch = big
                    bs.save(update_fields=["zt_batch", "update_time"])
                return big
            except Exception as e:
                with _print_lock:
                    print(f"{bs.batch} 归属解析失败: {e}", flush=True)
                return None
            finally:
                done[0] += 1
                if done[0] % 1000 == 0:
                    with _print_lock:
                        print(f"归属解析 {done[0]}/{total}", flush=True)

        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
            results = list(pool.map(resolve_one, items))

        bigs = []  # 保序去重
        seen = set()
        for big in results:
            if big and big not in seen:
                seen.add(big)
                bigs.append(big)

    if fresh and not skip_resolve:
        # 清理: 曾被错当成大批计算过直通字段、如今无任何子批指向的批,
        # 剥离其直通键并解除自指锁定, 交由后续大批计算重新归属
        cleaned = 0
        for bs in BatchSt.objects.filter(
            create_time__gte=since_dt, version=1, data__has_key="直通_良率",
        ).iterator():
            if bs.batch in seen:
                continue
            if BatchSt.objects.filter(zt_batch=bs.batch, version=1).exclude(id=bs.id).exists():
                continue
            data = bs.data
            for k in ["直通_子批次", "直通_总合格数", "直通_白料数", "直通_良率", "直通_口径"]:
                data.pop(k, None)
            bs.data = json.loads(json.dumps(data, cls=MyJSONEncoder))
            if bs.zt_batch == bs.batch:
                bs.zt_batch = None
                bs.save(update_fields=["data", "zt_batch", "update_time"])
            else:
                bs.save(update_fields=["data", "update_time"])
            cleaned += 1
            print(f"清理错误大批标记: {bs.batch}", flush=True)
        print(f"清理 {cleaned} 个", flush=True)

    print(f"共 {len(bigs)} 个大批, 开始计算直通良率", flush=True)
    done[0] = 0

    def cal_one(big):
        try:
            res = cal_zt_big(big, use_white_cache=not fresh) or {}
            done[0] += 1
            with _print_lock:
                print(f"[{done[0]}/{len(bigs)}] {big} 良率={res.get('良率')} 白料={res.get('白料数')} "
                      f"合格={res.get('总合格数')} 子批={res.get('子批数')} 口径={res.get('口径')}", flush=True)
        except Exception as e:
            done[0] += 1
            with _print_lock:
                print(f"[{done[0]}/{len(bigs)}] {big} 计算失败: {e}", flush=True)

    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
        list(pool.map(cal_one, bigs))
    print("done")


if __name__ == "__main__":
    fresh = "--fresh" in sys.argv
    skip_resolve = "--skip-resolve" in sys.argv
    workers = 8
    args = []
    for a in sys.argv[1:]:
        if a in ("--fresh", "--skip-resolve"):
            continue
        if a.startswith("--workers="):
            workers = int(a.split("=", 1)[1])
            continue
        args.append(a)
    main(since=args[0] if args else "2026-06-01", fresh=fresh, workers=workers, skip_resolve=skip_resolve)