# -*- coding: utf-8 -*-
"""修复分镜段：尺寸/风格短语、运镜词入段落、[Shot k] At 前缀、<d> 台词、同框<=3。"""
import json, sys, re, glob

SIZE_PHRASE = {
    "extreme-wide": "extreme wide shot",
    "wide": "wide shot",
    "medium": "medium shot",
    "close": "close-up",
    "extreme-close": "extreme close-up",
}
CAM_WORDS = ["static shot","push in","pull out","tracking shot","pan left","pan right",
             "tilt up","tilt down","pedestal up","pedestal down","arc shot","zoom in","zoom out",
             "shake slightly","shake strongly","pov","roll clockwise","roll counterclockwise"]
CAM_CANON = {
    "static shot": "Static Shot", "push in": "Push In", "pull out": "Pull Out",
    "tracking shot": "Tracking Shot", "pan left": "Pan Left", "pan right": "Pan Right",
    "tilt up": "Tilt Up", "tilt down": "Tilt Down", "pedestal up": "Pedestal Up",
    "pedestal down": "Pedestal Down", "arc shot": "Arc Shot", "zoom in": "Zoom In",
    "zoom out": "Zoom Out", "shake slightly": "Shake Slightly", "shake strongly": "Shake Strongly",
    "pov": "POV", "roll clockwise": "Roll Clockwise", "roll counterclockwise": "Roll Counterclockwise",
}

def norm_size(sz):
    sz = (sz or "").strip().lower().replace(" shot","").replace(" shot "," ")
    sz = sz.replace("close-up","close").replace("close up","close")
    sz = sz.replace("extreme close","extreme-close").replace("extreme wide","extreme-wide")
    if sz not in SIZE_PHRASE:
        for k,v in SIZE_PHRASE.items():
            if v in (sz or ""): return k
        return "medium"
    return sz

def repair_segment(seg):
    cuts = seg.get("cuts") or []
    for c in cuts:
        c["size"] = norm_size(c.get("size"))
        c["seconds"] = round(float(c.get("seconds") or 0), 2)
        if not isinstance(c.get("beats"), list) or len(c["beats"]) != 2:
            c["beats"] = [0, 0]
        else:
            c["beats"] = [int(c["beats"][0]), int(c["beats"][1])]
        if not isinstance(c.get("characters"), list): c["characters"] = []
        if not isinstance(c.get("props"), list): c["props"] = []
        cam = str(c.get("camera") or "").strip()
        cam_low = cam.lower()
        zh_map = {"固定机位":"Static Shot","推近":"Push In","拉远":"Pull Out","跟随":"Tracking Shot",
                  "跟拍":"Tracking Shot","摇镜":"Pan Left","俯拍":"Tilt Up","仰拍":"Tilt Down",
                  "环绕":"Arc Shot","特写推":"Push In","缓推":"Push In","缓拉":"Pull Out"}
        for k,v in zh_map.items():
            if k in cam: cam = v; cam_low = cam.lower(); break
        hit = None
        for w in CAM_WORDS:
            if w in cam_low:
                hit = w; break
        if hit is None:
            cam = "Static Shot"; cam_low = "static shot"; hit = "static shot"
        c["camera"] = CAM_CANON[hit]
        ph = SIZE_PHRASE[c["size"]]
        fr = str(c.get("frame") or "").strip()
        if ph not in fr.lower():
            for k2, v2 in SIZE_PHRASE.items():
                fr = re.sub(re.escape(v2), "", fr, flags=re.I)
            fr = (ph + " of " + fr).strip() if fr else ph
            fr = re.sub(r"\s+", " ", fr)
        # 清洗非英文（CJK 等），保底英文模板
        fr_clean = re.sub(r"[\u4e00-\u9fff\u3000-\u303f\uff00-\uffef]", "", fr)
        fr_clean = re.sub(r"\s+", " ", fr_clean).strip(" ,.")
        if len(fr_clean) < 15:
            fr_clean = ph + " of the scene described, cinematic film still"
        fr = fr_clean
        if "cinematic film still" not in fr.lower():
            fr = (fr.rstrip(".,") + ", cinematic film still").strip()
        c["frame"] = fr
        if len(c["characters"]) > 3:
            c["characters"] = c["characters"][:3]
    soundscape = seg.get("sound") or ""
    music = seg.get("music") or ""
    h3 = seg.get("h3Prompt") or ""
    for ln in h3.split("\n"):
        if ln.startswith("overall_soundscape:"): soundscape = ln[len("overall_soundscape:"):].strip()
        elif ln.startswith("non_diegetic_music:"): music = ln[len("non_diegetic_music:"):].strip()
    return seg, soundscape, music

def h3_cut_time(t):
    m = int(t // 60); s = int(t % 60); ms = int(round((t - int(t)) * 1000))
    return "%02d:%02d.%03d" % (m, s, ms)

def rebuild(seg, soundscape, music, beat_map):
    cuts = seg["cuts"]
    out = []
    starts = [0.0]
    t = 0.0
    for c in cuts:
        t += c["seconds"]
        starts.append(round(t * 10) / 10.0)  # 与 validate 的 cutStarts 一致：累计先取 1 位小数
    if len(cuts) <= 1:
        out.append("For the target video, at 0.00 seconds into the target video, <Picture 1> (from [Shot 1]) is fully referenced.")
    else:
        parts = ["Picture %d (from Shot %d) aligns with the %.2f-second mark of the target video" % (i+1, i+1, starts[i]) for i in range(len(cuts))]
        out.append("How the reference pictures align with the target video — " + "; ".join(parts) + ".")
    out.append("integrated_multimodal_description:")
    for i, c in enumerate(cuts):
        content = str(c.get("frame") or "").strip()
        cam_low = str(c["camera"]).lower()
        if not any(w in content.lower() for w in CAM_WORDS):
            content = (content.rstrip("., ") + " (camera: %s)" % cam_low)
        for b in beat_map.get(i, []):
            d = "<d>[Chinese] %s</d>" % b["text"]
            if d not in content:
                content = content.rstrip() + " " + d
        if i == 0:
            out.append("[Shot 1] " + content)
        else:
            out.append("[Shot %d] At %s, the camera cuts to <Picture %d>: %s" % (i+1, h3_cut_time(starts[i]), i+1, content))
    out.append("overall_soundscape: " + (soundscape or "ambient room tone and movement sounds"))
    out.append("non_diegetic_music: " + (music or "N/A"))
    seg["h3Prompt"] = "\n".join(out)
    return seg

def repair_ep(ep):
    segs = ep.get("segments") or []
    beats_by_scene = {}
    for sc in ep.get("seedScenes", []):
        for b in sc.get("beats", []):
            beats_by_scene[(sc["sceneIndex"], b["n"])] = b
    props_by_scene = {}
    for sc in ep.get("seedScenes", []):
        props_by_scene[sc["sceneIndex"]] = set(sc.get("props") or [])
    for seg in segs:
        seg, soundscape, music = repair_segment(seg)
        cuts = seg["cuts"]
        allowed_props = props_by_scene.get(seg.get("sceneIndex"), set())
        for c in cuts:
            c["props"] = [p for p in (c.get("props") or []) if p in allowed_props]
        beat_map = {}
        for ci, c in enumerate(cuts):
            a, b = c["beats"]
            lst = []
            for n in range(a, b+1):
                bb = beats_by_scene.get((seg.get("sceneIndex"), n))
                if bb and bb.get("kind") == "line":
                    lst.append(bb)
            beat_map[ci] = lst
        # 台词容纳：cut 秒数 >= 该切认领台词的 1 位小数和（与 validate 口径一致）
        for ci, c in enumerate(cuts):
            dlg = 0.0
            a, b = c["beats"]
            for n in range(a, b + 1):
                bb = beats_by_scene.get((seg.get("sceneIndex"), n))
                if bb and bb.get("kind") == "line":
                    dlg += float(bb.get("seconds") or 0)
            dlg = round(dlg * 10) / 10.0
            if dlg > float(c["seconds"]):
                c["seconds"] = round(dlg, 2)
        rebuild(seg, soundscape, music, beat_map)
        for c in cuts:
            if len(c["characters"]) > 3:
                c["characters"] = c["characters"][:3]
    return ep

if __name__ == "__main__":
    files = sys.argv[1:] or sorted(glob.glob("小说2/storyboard-work/ep-[0-9]*.json"))
    for f in files:
        if f.endswith("-seed.json"): continue
        ep = json.load(open(f, encoding="utf-8"))
        nseg = len(ep.get("segments") or [])
        if nseg == 0: continue
        repair_ep(ep)
        json.dump(ep, open(f, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
        print("repaired", f, "segs:", nseg)