from pathlib import Path
import argparse, math, json, subprocess, functools
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageFilter
from fontTools.ttLib import TTFont

P = Path(__file__).resolve().parent
SOURCE = next((root for root in P.parents if (root/'higgsfield/v4').is_dir()), Path('/var/workspace/documents/personal-JMvNPxCK/Work/Marketing/videos/sauna-explainer'))
V4 = SOURCE / 'higgsfield/v4'
A = SOURCE / 'assets'
BUILD = P / 'build'
PRE = P / 'previews'
for x in [BUILD, PRE, P/'assets']:
    x.mkdir(exist_ok=True, parents=True)
W,H,FPS = 1920,1080,24
INK = '#171814'
GREEN = '#003116'
MINT = '#93efa4'
PAPER = '#f1eadb'
CREAM = '#e7dfcc'
MUTED = '#626356'
SANS = '/usr/share/fonts/dejavu-sans-fonts/DejaVuSans.ttf'
BOLD = '/usr/share/fonts/dejavu-sans-fonts/DejaVuSans-Bold.ttf'
MONO = '/usr/share/fonts/dejavu-sans-mono-fonts/DejaVuSansMono.ttf'
ANTON = P/'assets/Anton.ttf'
if not ANTON.exists():
    for f in sorted(A.glob('anton-*.woff2')):
        font=TTFont(f)
        if all(ord(c) in font.getBestCmap() for c in 'WHAT IT CAN TOUCH'):
            font.flavor=None
            font.save(ANTON)
            break
assert ANTON.exists()

@functools.lru_cache(None)
def font(size, style='sans'):
    return ImageFont.truetype(str({'sans':SANS,'bold':BOLD,'mono':MONO,'display':ANTON}[style]),size)

def text(im, xy, value, size=34, color=INK, style='bold', center=False):
    d=ImageDraw.Draw(im)
    f=font(size,style)
    if center:
        width=d.textlength(value,font=f)
        xy=(xy[0]-width/2,xy[1])
    d.text(xy,value,font=f,fill=color,anchor='lt',stroke_width=0)

def texture(hexcolor):
    rgb=tuple(int(hexcolor[i:i+2],16) for i in (1,3,5))
    rng=np.random.default_rng(83)
    n=rng.normal(0,1.25,(H,W,1))
    a=np.clip(np.array(rgb)[None,None,:]+n,0,255).astype('uint8')
    return Image.fromarray(a,'RGB').convert('RGBA')
BACK=texture(CREAM)
DARK=texture(GREEN)

def paper(w,h,fill=PAPER,seed=5):
    im=Image.new('RGBA',(w+30,h+36))
    d=ImageDraw.Draw(im)
    rng=np.random.default_rng(seed)
    xs=list(range(8,w+8,28))
    pts=[(x,8+int(rng.integers(0,4))) for x in xs]+[(w+8,9),(w+8,h+8)]
    pts += [(x,h+8-int(rng.integers(0,4))) for x in reversed(xs)] + [(8,h+8)]
    shadow=Image.new('RGBA',im.size)
    ImageDraw.Draw(shadow).polygon([(x+6,y+8) for x,y in pts],fill=(15,25,15,62))
    im.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(5)))
    ImageDraw.Draw(im).polygon(pts,fill=fill)
    return im

def place(im, layer, x,y, opacity=1,angle=0):
    if opacity<=0:return
    q=layer
    if angle:q=q.rotate(angle,Image.Resampling.BICUBIC,expand=True)
    if opacity<1:
        q=q.copy();q.putalpha(q.getchannel('A').point(lambda p:int(p*opacity)))
    im.alpha_composite(q,(int(x-(q.width-layer.width)/2),int(y-(q.height-layer.height)/2)))

def enter(t,start,dur=.28):
    return 1-(1-max(0,min(1,(t-start)/dur)))**3

def heading(im,title,dark=False):
    size=100
    while font(size,'display').getlength(title)>1680:size-=2
    width=int(font(size,'display').getlength(title))+60
    strip=paper(width,148,INK if dark else PAPER,10)
    text(strip,(36,24),title,size,PAPER if dark else INK,'display')
    place(im,strip,85,48,angle=.6)

def dash(im,points,fill,width=5,dashlen=19,gap=15):
    d=ImageDraw.Draw(im)
    for a,b in zip(points,points[1:]):
        dx,dy=b[0]-a[0],b[1]-a[1]; L=math.hypot(dx,dy)
        for s in np.arange(0,L,dashlen+gap):
            e=min(L,s+dashlen)
            d.line([(a[0]+dx*s/L,a[1]+dy*s/L),(a[0]+dx*e/L,a[1]+dy*e/L)],fill=fill,width=width)

def path(im,points,color=INK,width=9,progress=1,head=True):
    d=ImageDraw.Draw(im)
    lengths=[math.dist(a,b) for a,b in zip(points,points[1:])]
    budget=sum(lengths)*progress
    done=[]
    for a,b,L in zip(points,points[1:],lengths):
        if budget<=0:break
        r=min(1,budget/L); e=(a[0]+(b[0]-a[0])*r,a[1]+(b[1]-a[1])*r)
        d.line([a,e],fill=color,width=width);done=[a,e];budget-=L
    if head and done:
        a,e=done;ang=math.atan2(e[1]-a[1],e[0]-a[0]);size=19
        d.polygon([e,(e[0]-size*math.cos(ang-.55),e[1]-size*math.sin(ang-.55)),(e[0]-size*math.cos(ang+.55),e[1]-size*math.sin(ang+.55))],fill=color)

def point_on(points,v):
    lengths=[math.dist(a,b) for a,b in zip(points,points[1:])];bgt=sum(lengths)*v
    for a,b,L in zip(points,points[1:],lengths):
        if bgt<=L:return(a[0]+(b[0]-a[0])*bgt/L,a[1]+(b[1]-a[1])*bgt/L)
        bgt-=L
    return points[-1]

@functools.lru_cache(None)
def face(name,size):
    src=Image.open(A/f'cut/raw_{name}.png').convert('RGB')
    arr=np.array(src).astype('int16');r,g,b=[arr[:,:,i] for i in range(3)]
    key=(g>r+10)&(g>b+8)
    alpha=np.where(key,0,255).astype('uint8')
    src=src.convert('RGBA');src.putalpha(Image.fromarray(alpha))
    src=src.crop(src.getbbox());src.thumbnail((size,size),Image.Resampling.LANCZOS)
    src.save(P/f'assets/{name}_{size}.png')
    return src

@functools.lru_cache(None)
def tile(label,w=245,h=120,fill=PAPER,small=None):
    p=paper(w,h,fill)
    fg=PAPER if fill==INK or fill==GREEN else INK
    sz=38
    while font(sz,'bold').getlength(label)>w-36:sz-=1
    text(p,(w/2+8,32),label,sz,fg,center=True)
    if small:text(p,(w/2+8,85),small,23,fg,'sans',True)
    return p


def schematic(t):
    im=BACK.copy();heading(im,'WHAT IT CAN TOUCH',True)
    dash(im,[(388,275),(1835,275),(1835,925),(388,925),(388,275)],MUTED,4)
    tag=tile("FILIP'S GUARDRAILS",465,74,INK)
    place(im,tag,1325,232)
    text(im,(455,965),'Illustration: access and actions stay within the permissions Filip grants.',25,MUTED,'sans')
    q=enter(t,.55)
    angle=3*math.sin((t-1.8)*12) if 1.8<t<2.4 else 0
    place(im,face('robert',225),80,400,q,angle)
    if q>.9:
        text(im,(192,655),'Robert',36,center=True)
        text(im,(192,315),'QUESTION',30,MUTED,'bold',True)
    p1=enter(t,1.7,.4)
    if p1:path(im,[(310,520),(448,520)],progress=p1)
    sa=paper(245,265,GREEN)
    text(sa,(130,54),"Filip's",42,PAPER,center=True)
    text(sa,(130,120),'Sauna',64,MINT,'display',True)
    text(sa,(130,219),'checks access',25,PAPER,'sans',True)
    place(im,sa,450,405,enter(t,2.1))
    p2=enter(t,2.7,.4)
    if p2:path(im,[(720,520),(803,520)],progress=p2)
    if t>2.8:text(im,(970,340),'MEMORY',30,MUTED,'bold',True)
    for k,s in enumerate(['USER_PROFILE','RULES','COMPANY']):
        place(im,tile(s,295,78),812,404+k*94,enter(t,2.9+k*.14))
    if t>4.1:text(im,(1309,340),'TOOLS',30,MUTED,'bold',True)
    if t>4.0:path(im,[(1134,520),(1180,520)],progress=enter(t,4.0,.25))
    for k,(s,cue) in enumerate([('PostHog',4.2),('Gmail',5.2),('Linear',5.7)]):
        place(im,tile(s,230,78,MINT),1190,404+k*94,enter(t,cue))
    if t>6.25:
        gate=tile('ONLY IF ALLOWED',320,65,INK)
        place(im,gate,1470,350,enter(t,6.25))
    if t>7.2:path(im,[(1445,520),(1528,520)],progress=enter(t,7.2,.35))
    act=paper(260,203,INK)
    text(act,(138,40),'ACTION',43,MINT,'display',True)
    text(act,(138,104),'within scope',26,PAPER,'sans',True)
    text(act,(138,150),'then report back',24,PAPER,'sans',True)
    place(im,act,1530,442,enter(t,7.55))
    ret=[(1665,673),(1665,800),(191,800),(191,715)]
    if t>8.65:
        path(im,ret,GREEN,8,enter(t,8.65,1.0))
        lab=tile('ANSWER + SOURCES',455,78,MINT)
        place(im,lab,665,755,enter(t,9.4))
        if 8.8<t<10.1:
            x,y=point_on(ret,max(0,min(1,(t-8.8)/1.3)))
            ImageDraw.Draw(im).ellipse((x-13,y-13,x+13,y+13),fill=MINT,outline=GREEN,width=3)
    if t>10.2:
        dash(im,[(388,275),(1835,275),(1835,925),(388,925),(388,275)],GREEN,6)
    place(im,tag,1325,232)
    return im.convert('RGB')


def receipt(t):
    im=BACK.copy();heading(im,'SOURCES ATTACHED',True)
    p=paper(1060,645)
    text(p,(60,43),'ILLUSTRATIVE ANSWER',27,MUTED,'mono')
    text(p,(60,101),'An answer you can check.',54,INK,'bold')
    ImageDraw.Draw(p).line([(62,198),(1010,198)],fill=INK,width=3)
    for k,(lab,val,cue) in enumerate([('READ','RULES.md',1.4),('CHECKED','PostHog',2.5),('FOLLOWED','The relevant access rules',3.5)]):
        if t>cue:
            y=239+k*94
            text(p,(64,y),lab,29,GREEN,'mono')
            text(p,(315,y),val,34,INK,'sans')
    d=ImageDraw.Draw(p)
    for x in range(10,1055,40):d.polygon([(x,650),(x+20,670),(x+40,650)],fill=PAPER)
    text(p,(62,569),'Open the source. Check the answer.',31,GREEN,'bold')
    place(im,p,410,276+25*(1-enter(t,.15)),enter(t,.1),angle=-.4)
    return im.convert('RGB')


def permissions(t):
    im=DARK.copy();heading(im,'PERMISSIONS BY RELATIONSHIP')
    for x,v in [(390,'PERSON'),(825,'ASK'),(1150,'ACT')]:text(im,(x,290),v,28,MINT)
    for k,(name,real,act) in enumerate([('Robert','robert',True),('Rafa','rafael',True),('Sebastian','sebastian',False),('Dasol','dasol',False)]):
        cue=.3+k*.18; q=enter(t,cue);y=351+k*136
        row=paper(1560,118,'#124026',seed=k+40)
        text(row,(220,44),name,38,PAPER)
        place(row,face(real,107),75,6)
        place(row,tile('ASK',195,68,PAPER),635,17)
        if act:place(row,tile('ACT',195,68,MINT),955,17,enter(t,3.2))
        else:
            dash(row,[(965,27),(1153,27),(1153,88),(965,88),(965,27)],'#a7b6a5',3,10,9)
        place(im,row,160-25*(1-q),y,q)
    text(im,(195,986),'Illustrated access levels. Permissions remain under the owner\'s control.',26,PAPER,'sans')
    return im.convert('RGB')


def run(cmd):
    r=subprocess.run([str(a) for a in cmd],capture_output=True,text=True)
    if r.returncode:raise RuntimeError(r.stderr[-3000:])
    return r

def probe(p):
    return json.loads(run(['ffprobe','-v','error','-show_streams','-show_format','-of','json',p]).stdout)

def duration(p):return float(probe(p)['format']['duration'])

def render_scene(key,fn,duration_s):
    dest=BUILD/f'{key}.mp4'
    count=round(duration_s*FPS)
    log=(BUILD/f'{key}_encode.log').open('w')
    enc=subprocess.Popen(['ffmpeg','-v','error','-y','-f','rawvideo','-pix_fmt','rgb24','-s','1920x1080','-r','12','-i','-','-an','-r','24','-frames:v',str(count),'-c:v','libx264','-threads','1','-preset','fast','-crf','18','-pix_fmt','yuv420p',str(dest)],stdin=subprocess.PIPE,stderr=log)
    for i in range(math.ceil(count/2)):
        enc.stdin.write(fn(i/12).tobytes())
    enc.stdin.close()
    if enc.wait():raise RuntimeError(f'encode failed {key}')
    print('RENDER_DONE',key,duration(dest),flush=True)


def assemble():
    order=[0,1,2,3,4,5,6,7,8,10,11]
    clip_names=['00_bottleneck','01_theyask_sauna','02_taste','03_theyask','04_receipts','05_perms','06_touch','08_rejected_v2','09_escalate','04_iask','10_beach']
    script=json.loads((V4/'script.json').read_text())
    script[3]='Dasol asks if he built an app for the off-site. Rafa asks which positioning he landed on.'
    timeline=[];start_frame=0
    for old in order:
        vo=V4/f'vo/vo_{old:02d}.mp3'
        take_len=4.76 if old==3 else duration(vo)
        lead=.30
        nframes=math.ceil((lead+take_len+.52)*FPS)
        if old==10:nframes=max(nframes,96)
        if old==11:nframes=max(nframes,96)
        length=nframes/FPS
        target=BUILD/f'seg_{old:02d}.mp4'
        custom={4:receipt,5:permissions,6:schematic}
        print('SEGMENT_START',old,length,flush=True)
        if old in custom:
            render_scene(f'seg_{old:02d}',custom[old],length)
            visual=str(target)
        elif old==11:
            src=SOURCE/'higgsfield/frames/11_end.png'
            run(['ffmpeg','-v','error','-y','-loop','1','-i',src,'-vf','scale=1920:1080,setsar=1','-frames:v',nframes,'-r',FPS,'-an','-c:v','libx264','-threads','1','-preset','fast','-crf','18','-pix_fmt','yuv420p',target])
            visual=str(src)
        else:
            src=V4/f'clips/{clip_names[old]}.mp4'
            run(['ffmpeg','-v','error','-y','-i',src,'-vf',f'fps=24,scale=1920:1080,setsar=1,tpad=stop_mode=clone:stop_duration=2,trim=duration={length},setpts=PTS-STARTPTS','-frames:v',nframes,'-an','-c:v','libx264','-threads','1','-preset','fast','-crf','18','-pix_fmt','yuv420p',target])
            visual=str(src)
        normalized=BUILD/f'vo_{old:02d}.wav'
        trim=f'atrim=0:{take_len},afade=t=out:st={take_len-.014}:d=0.014,' if old==3 else ''
        run(['ffmpeg','-v','error','-y','-i',vo,'-af',trim+'loudnorm=I=-18:TP=-3:LRA=8','-ar','48000','-ac','2','-c:a','pcm_s16le',normalized])
        timeline.append({'source_scene':old,'start_frame':start_frame,'frames':nframes,'start':start_frame/FPS,'duration':length,'vo_delay':lead,'vo_length':take_len,'voice_file':str(normalized),'segment':str(target),'visual_source':visual,'narration':script[old]})
        start_frame+=nframes
        print('SEGMENT_DONE',old,flush=True)
    total=start_frame/FPS
    (P/'timeline.json').write_text(json.dumps({'fps':FPS,'total_frames':start_frame,'duration':total,'removed_source_scene':9,'denial_buzzer':False,'scenes':timeline},indent=2))
    (P/'script.json').write_text(json.dumps([s['narration'] for s in timeline],indent=2))
    (BUILD/'concat.txt').write_text(''.join("file '"+s['segment']+"'\n" for s in timeline))
    run(['ffmpeg','-v','error','-y','-f','concat','-safe','0','-i',BUILD/'concat.txt','-c','copy',BUILD/'video.mp4'])
    finish_audio(timeline,total)


def finish_audio(timeline,total):
    inputs=[];filters=[]
    for i,s in enumerate(timeline):
        inputs+=['-i',s['voice_file']];d=round((s['start']+s['vo_delay'])*48000)
        filters.append(f'[{i}:a]adelay={d}S|{d}S[v{i}]')
    filters.append(''.join(f'[v{i}]' for i in range(len(timeline)))+f'amix=inputs={len(timeline)}:normalize=0:dropout_transition=0,apad,atrim=0:{total}[vo]')
    (BUILD/'voice_filter.txt').write_text(';'.join(filters))
    run(['ffmpeg','-v','error','-y']+inputs+['-filter_complex_script',BUILD/'voice_filter.txt','-map','[vo]','-c:a','pcm_s16le',P/'narration.wav'])
    run(['ffmpeg','-v','error','-y','-i',A/'bgm_long.mp3','-af',f'loudnorm=I=-35:TP=-9:LRA=8,atrim=0:{total},afade=t=in:d=1,afade=t=out:st={total-3.0}:d=3','-ar','48000','-ac','2',P/'music.wav'])
    idx={s['source_scene']:s for s in timeline}
    effect_inputs=['-i',P/'music.wav']
    effect_filters=[]
    for k,(effect,cue) in enumerate([('ding',idx[8]['start']+3.2),('clink',idx[10]['start']+.8)],1):
        effect_inputs+=['-i',A/f'sfx_{effect}.mp3']
        delay=round(cue*48000)
        effect_filters.append(f'[{k}:a]loudnorm=I=-30:TP=-9:LRA=7,volume=0.5,aresample=48000,aformat=channel_layouts=stereo,adelay={delay}S|{delay}S[e{k}]')
    effect_filters.append(f'[0:a][e1][e2]amix=inputs=3:normalize=0:dropout_transition=0,apad,atrim=0:{total}[bed]')
    (BUILD/'effects_filter.txt').write_text(';'.join(effect_filters))
    run(['ffmpeg','-v','error','-y']+effect_inputs+['-filter_complex_script',BUILD/'effects_filter.txt','-map','[bed]','-c:a','pcm_s16le',P/'music_effects.wav'])
    mix='[0:a]asplit=2[voice][key];[1:a][key]sidechaincompress=threshold=0.012:ratio=3:attack=20:release=240[quiet];[voice][quiet]amix=inputs=2:normalize=0:dropout_transition=0[mix]'
    run(['ffmpeg','-v','error','-y','-i',P/'narration.wav','-i',P/'music_effects.wav','-filter_complex',mix,'-map','[mix]','-c:a','pcm_s16le',BUILD/'premix.wav'])
    r=run(['ffmpeg','-hide_banner','-i',BUILD/'premix.wav','-af','loudnorm=I=-16:TP=-1.5:LRA=10:print_format=json','-f','null','-'])
    measured=json.JSONDecoder().raw_decode(r.stderr[r.stderr.rfind('{'):])[0]
    ln=f"loudnorm=I=-16:TP=-1.5:LRA=10:measured_I={measured['input_i']}:measured_TP={measured['input_tp']}:measured_LRA={measured['input_lra']}:measured_thresh={measured['input_thresh']}:offset={measured['target_offset']}:linear=true:print_format=json"
    run(['ffmpeg','-v','error','-y','-i',BUILD/'premix.wav','-af',ln,'-ar','48000','-ac','2',P/'mix.wav'])
    out=P/'there-while-not-there-v5.mp4'
    run(['ffmpeg','-v','error','-y','-i',BUILD/'video.mp4','-i',P/'mix.wav','-map','0:v:0','-map','1:a:0','-c:v','copy','-c:a','aac','-b:a','192k','-t',str(total),'-movflags','+faststart',out])
    checks={'output':probe(out),'removed_scene_9':all(s['source_scene']!=9 for s in timeline),'denial_effect_absent':True,'take_fits':[s['vo_delay']+s['vo_length']<=s['duration'] for s in timeline]}
    (P/'verification.json').write_text(json.dumps(checks,indent=2))
    print('FINAL',out,total,flush=True)
    thumbs=[]
    for i,s in enumerate(timeline):
        jpg=PRE/f'final_{i:02d}.jpg'
        run(['ffmpeg','-v','error','-y','-ss',str(s['start']+s['duration']*.68),'-i',out,'-frames:v','1','-vf','scale=640:360',jpg])
        thumb=Image.open(jpg).convert('RGB')
        d=ImageDraw.Draw(thumb);d.rectangle((0,0,640,28),fill=INK);d.text((12,5),f'{i+1:02d} | {s["start"]:.2f}s',font=font(18),fill='white');thumbs.append(thumb)
    contact=Image.new('RGB',(1920,360*math.ceil(len(thumbs)/3)),GREEN)
    for i,img in enumerate(thumbs):contact.paste(img,((i%3)*640,(i//3)*360))
    contact.save(P/'contact_sheet.jpg',quality=92)

if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('--render',action='store_true');parser.add_argument('--mix-only',action='store_true');args=parser.parse_args()
    for key,fn,times in [('schematic',schematic,[2.5,6.0,11.5]),('receipt',receipt,[4.7]),('permissions',permissions,[5.5])]:
        for t in times:fn(t).save(PRE/f'{key}_{t}.png')
    print('PREVIEWS_READY',flush=True)
    if args.render:assemble()
    if args.mix_only:
        data=json.loads((P/'timeline.json').read_text())
        finish_audio(data['scenes'],data['duration'])
