"""CPU専用子プロセス。Gradio・spaces・GPU予約を読み込まない。"""
import json,os,re,sys,uuid,time
from pathlib import Path

def atomic_json(path,data):
 temporary=path.with_name(path.name+'.'+uuid.uuid4().hex+'.tmp')
 temporary.write_text(json.dumps(data,ensure_ascii=False),encoding='utf-8');os.replace(temporary,path)

def run(root,jid):
 started=time.perf_counter()
 if not re.fullmatch(r'[a-f0-9]{32}',jid):raise ValueError('invalid_job')
 root=Path(root).resolve();directory=root/jid
 spec=json.loads((directory/'cpu-request.json').read_text(encoding='utf-8'))
 import torch
 torch.set_num_threads(1)
 from ultra.lunaron_model import LunaronBaseSeparator
 from audio_export import save_stem_wav
 import job_control
 def notify(progress,message):
  job_control.check_cancelled(root,jid)
  atomic_json(directory/'cpu-progress.json',{'progress':progress,'status':message})
 quality=spec['quality'];stems=spec['stems']
 if stems not in (2,4,6) or quality not in ('standard','hq','ultra'):raise ValueError('invalid_settings')
 path=Path(spec['input_path'])
 if spec.get('is_yt'):
  import yt_dlp
  with yt_dlp.YoutubeDL({'format':'bestaudio/best','outtmpl':str(directory/'input.%(ext)s'),'postprocessors':[{'key':'FFmpegExtractAudio','preferredcodec':'wav','preferredquality':'0'}],'quiet':True,'no_warnings':True,'noplaylist':True}) as downloader:downloader.download([spec['input_path']])
  path=directory/'input.wav'
 if not path.is_file():raise FileNotFoundError('input_missing')
 if quality=='ultra':
  from ultra.inference_policy import ultra_settings
  shifts,overlap=ultra_settings(stems);model_stems=stems
 else:
  shifts=(4 if stems==6 else 10) if quality=='hq' else 2
  overlap=.50;model_stems=6 if stems==6 else 4
 notify(0,'CPUで分離モデルを準備中です。GPU処理より時間がかかります。')
 separator=LunaronBaseSeparator(model_name='htdemucs_6s' if stems==6 else 'htdemucs_ft',stems_count=model_stems,device='cpu',shifts=shifts,overlap=overlap,callback=lambda d:notify(round(max(0,min(1,d.get('progress',0)))*80,1),'CPUで音声を分離しています。'),preserve_model_output=quality=='ultra')
 prepared=time.perf_counter()
 with job_control.inference_scope(root,jid):origin,parts=separator.separate_audio_file(str(path))
 inferred=time.perf_counter()
 if quality=='ultra' and origin is not None and parts and all(v is not None for v in parts.values()):
  length=min(origin.shape[-1],*(v.shape[-1] for v in parts.values()))
  parts={name:value[..., :length].cpu() for name,value in parts.items()}
 if stems==2 and quality!='ultra':
  accompaniment=None
  for name,tensor in parts.items():
   if name!='vocals':accompaniment=tensor.clone() if accompaniment is None else accompaniment+tensor
  parts={'vocals':parts.get('vocals'),'accompaniment':accompaniment}
 if len(parts)!=stems or any(v is None for v in parts.values()):raise RuntimeError('incomplete_stems')
 for index,(name,tensor) in enumerate(parts.items(),1):
  if name not in ('vocals','accompaniment','drums','bass','other','guitar','piano'):raise ValueError('invalid_stem')
  job_control.check_cancelled(root,jid)
  save_stem_wav(tensor,str(directory/(name+'.wav')),samplerate=separator.samplerate)
  notify(round(80+index/len(parts)*19,1),'CPU分離のトラックを保存しています。')
 atomic_json(directory/'cpu-result.json',{'stems':list(parts),'input_path':str(path),'timings_ms':{'prepare':round((prepared-started)*1000),'inference':round((inferred-prepared)*1000),'export':round((time.perf_counter()-inferred)*1000)}})

if __name__=='__main__':
 try:run(sys.argv[1],sys.argv[2])
 except Exception as error:
  # 自由文・音源名・パスを公開エラーに含めない。
  try:atomic_json(Path(sys.argv[1])/sys.argv[2]/'cpu-error.json',{'error_class':type(error).__name__})
  except Exception:pass
  sys.exit(1)