#!/usr/bin/env python3
import json, re, subprocess, time
from pathlib import Path

PROJECT=Path('/data/video-pipeline/Youtube-Video-Maker/Project/01062026-004-Ba-Trieu')
MASTER=PROJECT/'audio/master_narration.wav'
TIMELINE=PROJECT/'audio/master_narration_timeline.json'
SUBS=PROJECT/'subtitles/master_subtitles_clamped.srt'
PROBE=PROJECT/'audio/master_audio_probe.json'
PARTS=PROJECT/'parts'
AUDIO_DIR=PARTS/'audio'
SUB_DIR=PARTS/'subtitles'
RANGES=PARTS/'part_ranges.json'
VALIDATION=PARTS/'part_validation.json'
MANIFEST=PROJECT/'logs/step13_split_parts_manifest.json'
TARGET_PART_SEC=150.0
MIN_PART_SEC=90.0
SAFE_GAP=0.02
SSH=['sshpass','-p','11210119','ssh','-o','StrictHostKeyChecking=no','-o','PreferredAuthentications=password','-o','PubkeyAuthentication=no','root@192.168.1.104']

def log(m): print(time.strftime('%Y-%m-%d %H:%M:%S ')+m, flush=True)

def parse_t(v):
 m=re.match(r'^(\d+):(\d+):(\d+),(\d+)$',v.strip());
 if not m: raise ValueError(v)
 h,mi,s,ms=map(int,m.groups()); return h*3600+mi*60+s+ms/1000

def fmt_t(sec):
 sec=max(0,sec); ms=int(round(sec*1000)); h,rem=divmod(ms,3600000); mi,rem=divmod(rem,60000); s,ms=divmod(rem,1000); return f'{h:02d}:{mi:02d}:{s:02d},{ms:03d}'

def read_srt(p):
 txt=p.read_text(encoding='utf-8-sig').strip(); cues=[]
 if not txt: return cues
 for block in re.split(r'\n\s*\n',txt):
  lines=[x.rstrip('\r') for x in block.splitlines() if x.strip()]
  if len(lines)<2: continue
  tl=lines[1] if '-->' in lines[1] else lines[0]
  text='\n'.join(lines[2:] if '-->' in lines[1] else lines[1:]).strip()
  a,b=[x.strip() for x in tl.split('-->',1)]
  cues.append({'start':parse_t(a),'end':parse_t(b),'text':text})
 return cues

def write_srt(cues,p):
 out=[]
 for i,c in enumerate(cues,1):
  out += [str(i), f"{fmt_t(c['start'])} --> {fmt_t(c['end'])}", c['text'], '']
 p.write_text('\n'.join(out), encoding='utf-8')

def choose_ranges(timeline,total):
 scenes=timeline['timeline']; ranges=[]; start=0.0; start_idx=0
 while start_idx < len(scenes):
  target=start+TARGET_PART_SEC
  if total-target < MIN_PART_SEC and total-start <= TARGET_PART_SEC+MIN_PART_SEC:
   end=total; end_idx=len(scenes)-1
  else:
   candidates=[(i,s) for i,s in enumerate(scenes[start_idx:], start_idx) if s['end']>=start+MIN_PART_SEC and s['end']<=target+35]
   if candidates:
    # Prefer heading boundary just before a new section, otherwise closest to target.
    heading_prev=[(i,s) for i,s in candidates if i+1 < len(scenes) and scenes[i+1].get('scene_type')=='heading']
    pool=heading_prev or candidates
    end_idx,end_scene=min(pool, key=lambda x: abs(x[1]['end']-target))
    end=end_scene['end']
   else:
    end_idx=min(range(start_idx,len(scenes)), key=lambda i: abs(scenes[i]['end']-target)); end=scenes[end_idx]['end']
  part_scenes=scenes[start_idx:end_idx+1]
  ranges.append({'part_id':f'part_{len(ranges)+1:02d}','start_sec':round(start,3),'end_sec':round(end,3),'duration_sec':round(end-start,3),'scene_ids':[s['scene_id'] for s in part_scenes]})
  start=end; start_idx=end_idx+1
 return ranges

def cut_audio(start,dur,out):
 cmd=SSH+[f"ffmpeg -y -hide_banner -loglevel error -ss {start:.3f} -t {dur:.3f} -i {str(MASTER)!r} -ac 1 -ar 22050 {str(out)!r}"]
 subprocess.run(cmd, check=True)

def probe_audio(p):
 cmd=f"curl -sS -X POST http://192.168.1.104:8002/probe -H 'Content-Type: application/json' -d '{{\"input_path\":\"{p}\"}}'"
 res=subprocess.check_output(cmd, shell=True, text=True)
 doc=json.loads(res); return float(doc.get('format',{}).get('duration') or 0)

def main():
 AUDIO_DIR.mkdir(parents=True,exist_ok=True); SUB_DIR.mkdir(parents=True,exist_ok=True); MANIFEST.parent.mkdir(parents=True,exist_ok=True)
 timeline=json.loads(TIMELINE.read_text(encoding='utf-8'))
 total=float(json.loads(PROBE.read_text(encoding='utf-8'))['master_duration_sec'])
 cues=read_srt(SUBS)
 ranges=choose_ranges(timeline,total)
 validations=[]
 for r in ranges:
  pid=r['part_id']; audio=AUDIO_DIR/f'{pid}.wav'; sub=SUB_DIR/f'{pid}.srt'
  r['audio_path']=str(audio.relative_to(PROJECT)); r['subtitle_path']=str(sub.relative_to(PROJECT))
  log(f'CUT {pid} {r["start_sec"]}->{r["end_sec"]}')
  cut_audio(r['start_sec'], r['duration_sec'], audio)
  part_cues=[]
  for c in cues:
   if c['end'] <= r['start_sec'] or c['start'] >= r['end_sec']: continue
   ns=max(0.0,c['start']-r['start_sec']); ne=min(r['duration_sec'],c['end']-r['start_sec'])
   if ne <= ns: continue
   part_cues.append({'start':round(ns,3),'end':round(ne,3),'text':c['text']})
  # ensure no overlaps after rebase/clamp
  overlaps=0; invalid=0
  for i,c in enumerate(part_cues):
   if i and c['start'] < part_cues[i-1]['end']:
    part_cues[i-1]['end']=max(part_cues[i-1]['start']+0.2,c['start']-SAFE_GAP)
   if c['end']<=c['start'] or not c['text'].strip(): invalid+=1
  write_srt(part_cues, sub)
  real=probe_audio(audio); diff=real-r['duration_sec']
  for i,c in enumerate(part_cues):
   if i and c['start'] < part_cues[i-1]['end']: overlaps+=1
   if c['start']<0 or c['end']>real+0.05: invalid+=1
  validations.append({'part_id':pid,'expected_duration_sec':r['duration_sec'],'audio_duration_sec':real,'duration_diff_sec':round(diff,3),'subtitle_cues':len(part_cues),'overlap_remaining':overlaps,'invalid_cues':invalid})
 RANGES.write_text(json.dumps({'project_id':PROJECT.name,'step':13,'source_master_audio':str(MASTER),'source_subtitle':str(SUBS),'parts':ranges},ensure_ascii=False,indent=2),encoding='utf-8')
 status='completed' if all(v['overlap_remaining']==0 and v['invalid_cues']==0 and abs(v['duration_diff_sec'])<0.08 for v in validations) else 'failed_validation'
 VALIDATION.write_text(json.dumps({'project_id':PROJECT.name,'step':13,'status':status,'parts_count':len(ranges),'validations':validations},ensure_ascii=False,indent=2),encoding='utf-8')
 MANIFEST.write_text(json.dumps({'project_id':PROJECT.name,'step':13,'status':status,'part_ranges':str(RANGES),'validation':str(VALIDATION),'parts_count':len(ranges)},ensure_ascii=False,indent=2),encoding='utf-8')
 if status!='completed': raise SystemExit(status)
 log('DONE '+json.dumps({'status':status,'parts':len(ranges)},ensure_ascii=False))
if __name__=='__main__': main()