import json, unicodedata, string, argparse
from difflib import SequenceMatcher
import numpy as np
from moviepy.editor import VideoFileClip, AudioFileClip
from pydub import AudioSegment

FPS = 30
RATIO_MIN, RATIO_MAX = 0.3, 3.0
def crop_song(song_path, end):
    
    song = AudioSegment.from_file(song_path)
    end = min(song.duration_seconds, end + 0.5)
    song[:1000*end].export(song_path)
def _flatten_words(obj):
    words = []
    for w in obj["aligned_words"]:
            words.append({"raw": w["word"], "start": w["start_s"], "end": w["end_s"]})
    return words

def _norm(txt: str) -> str:
    tbl = str.maketrans('', '', string.punctuation + "‘’“”\"…—–")
    return unicodedata.normalize("NFKD", txt).translate(tbl).lower().strip()

def build_combined_sync(video_json, song_json):
    v_words = _flatten_words(video_json)
    s_words = _flatten_words(song_json)

    for lst in (v_words, s_words):
        for d in lst:
            d["norm"] = _norm(d["raw"])

    a = [w["norm"] for w in v_words]
    b = [w["norm"] for w in s_words]

    sm = SequenceMatcher(a=a, b=b, autojunk=False)
    print(sm)
    anchors = []
    prev_vs = prev_ss = -1.0

    for tag, i1, i2, j1, j2 in sm.get_opcodes():
        if tag != "equal":
            continue
        for vi, sj in zip(range(i1, i2), range(j1, j2)):
            vw, sw = v_words[vi], s_words[sj]
            vs0, ve0 = vw["start"], vw["end"]
            ss0, se0 = sw["start"], sw["end"]
            if vs0 <= prev_vs or ss0 <= prev_ss:
                continue
            ratio = (se0 - ss0) / max(ve0 - vs0, 1e-6)
            if RATIO_MIN <= ratio <= RATIO_MAX:
                anchors.append({
                    "songStart": ss0,
                    "songEnd":   se0,
                    "videoStart":vs0,
                    "videoEnd":  ve0
                })
                prev_vs, prev_ss = vs0, ss0
    if not anchors:
        raise RuntimeError("No anchors aligned – check transcripts.")
    return anchors

def smooth_monotonic(arr: np.ndarray, window: int) -> np.ndarray:
    arr = np.maximum.accumulate(arr)
    if window > 1 and len(arr) > window:
        kernel = np.ones(window) / window
        pad = window // 2
        padded = np.pad(arr, (pad, pad), mode='edge')
        arr = np.convolve(padded, kernel, mode='same')[pad:-pad]
    return arr

def retime_video(anchors, video_path, song_path, out_path, smooth_win):
    out_times, in_times = [], []
    for a in anchors:
        out_times += [a["songStart"], a["songEnd"]]
        in_times  += [a["videoStart"], a["videoEnd"]]
    video = VideoFileClip(video_path)
    last_video_anchor = in_times[-1]
    last_song_anchor = out_times[-1]
    print(f"last video anchor {last_video_anchor}")
    print(f"last song anchor {last_song_anchor}")
    print(video.duration)
    last_segement_dur = video.duration - last_video_anchor
    in_times += [last_video_anchor, video.duration]
    out_times += [last_song_anchor, last_song_anchor + last_segement_dur]
    out_t = np.array(out_times); in_t = np.array(in_times)

    idx = np.argsort(out_t)
    out_t, in_t = out_t[idx], in_t[idx]
    mask = np.append([True], np.diff(out_t) > 1e-6)
    out_t, in_t = out_t[mask], in_t[mask]

    in_t = smooth_monotonic(in_t, smooth_win)

    song  = AudioFileClip(song_path).set_duration(last_song_anchor + last_segement_dur)
    video = VideoFileClip(video_path)


    if out_t[0] > 0:
        out_t = np.insert(out_t, 0, 0.0)
        in_t  = np.insert(in_t,  0, in_t[0])
    # if out_t[-1] < song.duration:
    #     out_t = np.append(out_t, song.duration)
    #     in_t  = np.append(in_t, in_t[-1])

    def map_time(t):
        return float(np.interp(t, out_t, in_t))

    result = (video.fl_time(map_time)
                   .set_duration(last_song_anchor + last_segement_dur)
                   .set_audio(song))
    result.write_videofile(out_path,
                           fps=FPS,
                           codec="libx264",
                           audio_codec="aac",
                           threads=4)

def make_synced_video(clip_path, song_path, video_json_path, song_json_path,
                      out_path, export_sync=None, smooth=3):
    with open(video_json_path) as f: v_js = json.load(f)
    with open(song_json_path)  as f: s_js = json.load(f)

    anchors = build_combined_sync(v_js, s_js)

    if export_sync:
        with open(export_sync, "w") as f:
            json.dump(anchors, f, indent=2)

    retime_video(anchors, clip_path, song_path, out_path, smooth)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--clip", required=True)
    ap.add_argument("--song", required=True)
    ap.add_argument("--video_json", required=True)
    ap.add_argument("--song_json",  required=True)
    ap.add_argument("--out", default="synced.mp4")
    ap.add_argument("--smooth", type=int, default=3)
    ap.add_argument("--export_sync")
    args = ap.parse_args()

    make_synced_video(args.clip, args.song, args.video_json, args.song_json,
                      args.out, args.export_sync, args.smooth)

if __name__ == "__main__":
    main()