#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
offline_bg.pipeline — 五步编排：估光 → 重打光 → 闭环校验 → 画质降级 → 落库。

对外只有一个入口：
    prep = BackgroundPreprocessor(cfg)
    profile = profile_camera(frames)                 # ①（会话开始一次）
    rec = prep.process(bg, profile, user_id="u1")    # ②③④⑤（每张背景一次）

设计要点：
- 重打光在 work_size(768×448) 上做，画质匹配在**成片分辨率**上做（模糊 σ 随分辨率变）。
- ③ 不通过就换 strength + 换 seed 重来；全败也落库，但标 validated=false，
  由上层决定要不要用——不静默把失败品当成功品。
- IC-Light 模型懒加载：只做画质匹配（不重打光）时完全不碰 GPU 大模型。
"""
import time
from typing import Dict, List, Optional

import cv2
import numpy as np

from . import validated as V
from .config import DEFAULT, OfflineBGConfig
from .library import BackgroundLibrary
from .lighting import CameraProfile, profile_camera  # noqa: F401  (对外再导出)
from .quality import match_quality, quality_report
from .validate import best_of, score_direction


class BackgroundPreprocessor:
    def __init__(self, cfg: OfflineBGConfig = DEFAULT,
                 library: Optional[BackgroundLibrary] = None,
                 relighter=None, verbose: bool = True):
        self.cfg = cfg
        self.lib = library if library is not None else BackgroundLibrary(cfg.library_root, cfg)
        self._relighter = relighter
        self.verbose = verbose

    def _log(self, *a):
        if self.verbose:
            print("[offline_bg]", *a, flush=True)

    @property
    def relighter(self):
        """懒加载：不做重打光就不加载 SD（省 30–60s 与显存）。"""
        if self._relighter is None:
            from .iclight import ICLightRelighter
            t = time.perf_counter()
            self._relighter = ICLightRelighter.shared(self.cfg)
            self._log(f"IC-Light 就绪 {time.perf_counter()-t:.1f}s")
        return self._relighter

    # ------------------------------------------------------------------
    def _relight_with_retry(self, src_small: np.ndarray, side: str,
                            prompt_extra: str = "") -> Dict:
        """② + ③：重打光并闭环校验，不合格就换参数重来。"""
        c = self.cfg
        prompt = c.prompt_for(side, prompt_extra)
        outs, attempts = [], []
        for i in range(c.max_attempts):
            strength = c.retry_strengths[min(i, len(c.retry_strengths) - 1)]
            seed = c.base_seed + i * 7919          # 换 seed 才是真的换一次采样
            t = time.perf_counter()
            out = self.relighter.relight(src_small, side, prompt=prompt,
                                         strength=strength, seed=seed)
            sec = time.perf_counter() - t
            sc = score_direction(src_small, out, side, c)
            outs.append(out)
            attempts.append({"i": i, "strength": strength, "seed": seed,
                             "sec": round(sec, 2), "score": sc})
            self._log(f"  重打光#{i} strength={strength} {sec:.1f}s → "
                      f"{sc['axis']}差 {sc['src_diff']:+.1f}→{sc['out_diff']:+.1f} "
                      f"增量{sc['signed_delta']:+.1f} {'✔' if sc['ok'] else '✗ ' + sc['reason']}")
            if sc["ok"]:
                return {"image": out, "attempts": attempts, "score": sc,
                        "validated": True, "used": i}
        # 全败：挑最接近的一次留下，标 validated=false
        bi = best_of([a["score"] for a in attempts])
        self._log(f"  ⚠️ {c.max_attempts} 次均未达标，取最接近的 #{bi} 落库（validated=false）")
        return {"image": outs[bi], "attempts": attempts, "score": attempts[bi]["score"],
                "validated": False, "used": bi}

    # ------------------------------------------------------------------
    def process(self, bg_bgr: np.ndarray, profile: CameraProfile,
                user_id: str = "default", name: Optional[str] = None,
                prompt_extra: str = "", relight: bool = True,
                bg_id: Optional[str] = None) -> Dict:
        """把一张背景加工成可直接用于实时合成的成品并落库，返回完整记录。"""
        c = self.cfg
        t_all = time.perf_counter()
        side = profile.cls
        tw, th = c.target_size or profile.frame_size or (bg_bgr.shape[1], bg_bgr.shape[0])

        rec: Dict = {
            "user_id": str(user_id), "name": name,
            "camera_profile": profile.to_dict(),
            "target_light": {"cls": side, "dx": round(profile.dx, 4),
                             "dy": round(profile.dy, 4), "mag": round(profile.mag, 4)},
            "target_size": [tw, th],
            "source_impl": V.SOURCE_PATH,
            "config": {"work_size": list(c.work_size), "steps": c.steps,
                       "retry_strengths": list(c.retry_strengths), "cfg_scale": c.cfg_scale,
                       "min_delta": c.min_delta, "noise_gain": c.noise_gain,
                       "front_strategy": c.front_strategy},
        }

        # ---------- ②③ 重打光 + 闭环校验 ----------
        skip_reason = ""
        if not relight:
            skip_reason = "调用方指定 relight=False"
        elif side == "front" and c.front_strategy == "skip":
            skip_reason = "未检出明显方向光（front），按 front_strategy=skip 不重打光"
        elif not profile.light_reliable and side == "front":
            skip_reason = "光照方向不可靠（未检到人脸或方向不显著）"

        t_r = time.perf_counter()
        if skip_reason:
            self._log(f"跳过重打光：{skip_reason}")
            relit_full = cv2.resize(bg_bgr, (tw, th), interpolation=cv2.INTER_CUBIC) \
                if bg_bgr.shape[:2] != (th, tw) else bg_bgr.copy()
            rec.update({"relit": False, "relit_reason": skip_reason,
                        "attempts": [], "direction_score": None, "validated": False})
            relit_small = None
        else:
            src_small = cv2.resize(bg_bgr, c.work_size, interpolation=cv2.INTER_AREA)
            r = self._relight_with_retry(src_small, side, prompt_extra)
            relit_small = r["image"]
            relit_full = cv2.resize(relit_small, (tw, th), interpolation=cv2.INTER_CUBIC)
            rec.update({"relit": True, "relit_reason": "",
                        "attempts": r["attempts"], "direction_score": r["score"],
                        "validated": r["validated"], "attempt_used": r["used"]})
        relight_sec = time.perf_counter() - t_r

        # ---------- ④ 画质降级匹配（必须在成片分辨率上做）----------
        t_q = time.perf_counter()
        final, qdiag = match_quality(relit_full, profile, c, seed=c.base_seed,
                                     target_size=(tw, th))
        quality_sec = time.perf_counter() - t_q
        rec["quality"] = qdiag
        rec["quality_report"] = quality_report(relit_full, final, profile)
        if qdiag["warn"]:
            self._log(f"  ⚠️ {qdiag['warn']}")
        self._log(f"  画质匹配 {quality_sec:.2f}s 锐度 "
                  f"{rec['quality_report']['bg_sharp_before']} → "
                  f"{rec['quality_report']['bg_sharp_after']} "
                  f"(摄像头 {rec['quality_report']['camera_sharp']}, "
                  f"误差 {qdiag['sharp_err']*100:.1f}%)")

        rec["timing"] = {"relight_sec": round(relight_sec, 2),
                         "quality_sec": round(quality_sec, 2),
                         "total_sec": round(time.perf_counter() - t_all, 2)}

        # ---------- ⑤ 落库 ----------
        images = {"final": final, "original": bg_bgr}
        if relit_small is not None:
            images["relit"] = relit_full
        entry = self.lib.add(user_id, images, rec, name=name, bg_id=bg_id)
        rec["entry"] = entry
        self._log(f"已落库 bg_id={entry['bg_id']} cls={entry['cls']} "
                  f"validated={entry['validated']} 总耗时 {rec['timing']['total_sec']}s")
        return rec

    # ------------------------------------------------------------------
    def process_batch(self, bgs: List[np.ndarray], profile: CameraProfile,
                      user_id: str = "default", names: Optional[List[str]] = None,
                      **kw) -> List[Dict]:
        """批量：模型只加载一次，逐张过完整五步。某张炸掉不影响其余。"""
        out = []
        for i, bg in enumerate(bgs):
            nm = names[i] if names and i < len(names) else None
            try:
                out.append(self.process(bg, profile, user_id=user_id, name=nm, **kw))
            except Exception as e:                       # 单张失败不拖垮整批
                self._log(f"❌ 第 {i} 张失败: {type(e).__name__}: {e}")
                out.append({"error": f"{type(e).__name__}: {e}", "name": nm, "index": i})
        return out
