"""Bounded WAV tests with coordinated GPU ownership and disposable input."""
import io,json,os,signal,subprocess,threading,time,uuid,wave
from pathlib import Path
from image_test import probe
from separator_runtime import MODELS
class SeparatorTests:
 def __init__(self,root,runtime,scheduler):
  self.root=Path(root);self.runtime=runtime;self.scheduler=scheduler;self.lock=threading.RLock();self.job=None;self.process=None;self.cancel=threading.Event();self.thread=None
 def status(self):
  with self.lock:return dict(job=dict(self.job) if self.job else None)
 def start(self,model,audio,target=None,full=False):
  if full and self.scheduler.gpu_mode!='separator':raise ValueError('Bitte unter Steuerung Audio-Trennung starten.')
  targets={'vocals':'model_bs_roformer_ep_317_sdr_12.9755.ckpt','drums':'htdemucs_ft.yaml','bass':'htdemucs_ft.yaml','guitar':'htdemucs_6s.yaml','piano':'htdemucs_6s.yaml','other':'htdemucs_6s.yaml'}
  if target is not None and targets.get(target)!=model:raise ValueError('Zielspur und Modellpaket passen nicht zusammen.')
  if not isinstance(model,str) or model not in MODELS or not self.runtime.status()['installed'] or not self.runtime.model_ready(model):raise ValueError('Laufzeit und Modellpaket zuerst unter Audio Separator einrichten.')
  if not isinstance(audio,bytes) or not 44<=len(audio)<=(256 if full else 8)*1024**2:raise ValueError('WAV überschreitet die erlaubte Upload-Größe.')
  try:
   with wave.open(io.BytesIO(audio)) as wav:
    if wav.getnchannels() not in (1,2) or wav.getsampwidth()!=2 or not 8000<=wav.getframerate()<=48000 or not 0<wav.getnframes()/wav.getframerate()<=(1200 if full else 30):raise ValueError('PCM-WAV, 16 Bit, Mono/Stereo, maximal 30 Sekunden erforderlich.')
    if len(wav.readframes(wav.getnframes()))!=wav.getnframes()*wav.getnchannels()*2:raise ValueError('WAV-Datei ist unvollständig.')
  except (wave.Error,EOFError):raise ValueError('Ungültige WAV-Datei.') from None
  with self.lock:
   if self.job and self.job['state']=='running':raise ValueError('Trenntest läuft bereits.')
   self.target=target;self.full=full;self.cancel.clear();self.root.mkdir(parents=True,exist_ok=True,mode=0o700);self.job=dict(id=uuid.uuid4().hex,state='running',phase='Wartet auf GPU-Freigabe',started_at=time.time(),files=[])
   self.thread=threading.Thread(target=self._run,args=(model,audio),daemon=True);self.thread.start();return dict(self.job)
 def _run(self,model,audio):
  release=None;directory=self.root/self.job['id'];state='failed';phase='Trennung fehlgeschlagen.'
  try:
   release=self.scheduler.image_reservation(kind='separator')
   gpu=next((g for term in ('RTX 3060','RTX 5080') for g in probe() if term in g['name'] and not g['processes'] and g['free_mib']>4000),None)
   if not gpu:raise ValueError('Keine ausreichend freie GPU. Fremde Dienste werden nicht beendet.')
   directory.mkdir(mode=0o700);(directory/'input.wav').write_bytes(audio);diagnostics=(directory/'worker.tmp').open('wb')
   with self.lock:
    if self.cancel.is_set():raise InterruptedError()
    self.job.update(phase='Modell laden und Audio trennen',gpu=gpu['name'])
    self.process=subprocess.Popen([str(self.runtime.paths()[0]),str(Path(__file__).with_name('separator_worker.py')),str(self.runtime.root/'models'),model,str(directory/'input.wav'),str(directory),getattr(self,'target',None) or ''],stdout=subprocess.DEVNULL,stderr=diagnostics,start_new_session=True,env=dict(os.environ,CUDA_VISIBLE_DEVICES=gpu['uuid'],OMP_NUM_THREADS='6'))
   # Drain diagnostics into a bounded temporary file rather than retaining user data.
   deadline=time.monotonic()+(7200 if getattr(self,'full',False) else 600)
   while self.process.poll() is None:
    if self.cancel.wait(.2):raise InterruptedError()
    if time.monotonic()>deadline:raise ValueError('Trennauftrag überschreitet das Zeitlimit.')
    with self.lock:self.job['elapsed_seconds']=round(time.time()-self.job['started_at'])
   diagnostics.close();error=(directory/'worker.tmp').read_bytes()[-65536:].decode(errors='replace')
   if self.process.returncode:raise ValueError('GPU-Speicher reicht nicht (OOM).' if 'out of memory' in error.lower() else 'Audio Separator hat den Test abgebrochen. Modellpaket und Laufzeit prüfen.')
   names=json.loads((directory/'results.json').read_text())
   if not names or any(Path(n).name!=n or not n.endswith('.wav') or not (directory/n).is_file() for n in names):raise ValueError('Ungültige Ergebnisdateien.')
   with self.lock:self.job['files']=[dict(index=i,name=n) for i,n in enumerate(names)]
   (directory/'complete').touch();state='complete';phase='Spuren fertig · Modell entladen'
  except InterruptedError:state='cancelled';phase='Trenntest abgebrochen'
  except Exception as exc:phase=str(exc) if isinstance(exc,ValueError) else 'Trenntest fehlgeschlagen.'
  finally:
   process=self.process
   if process and process.poll() is None:
    try:os.killpg(process.pid,signal.SIGTERM);process.wait(5)
    except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
    except ProcessLookupError:pass
   if 'diagnostics' in locals():diagnostics.close()
   (directory/'input.wav').unlink(missing_ok=True);(directory/'worker.tmp').unlink(missing_ok=True)
   with self.lock:self.process=None;self.job.update(state=state,phase=phase,finished_at=time.time())
   if release:release()
 def stop(self):
  self.cancel.set()
  if self.thread:self.thread.join(7)
  return {'cancellation_requested':True}
 def audio(self,ident,index):
  if not isinstance(ident,str) or len(ident)!=32 or any(c not in '0123456789abcdef' for c in ident):raise ValueError('Ungültige Auftrag-ID.')
  directory=self.root/ident
  if not (directory/'complete').is_file():raise ValueError('Auftrag nicht fertig.')
  names=json.loads((directory/'results.json').read_text())
  if type(index) is not int or not 0<=index<len(names) or Path(names[index]).name!=names[index]:raise ValueError('Ungültige Spur.')
  path=directory/names[index]
  if path.is_symlink() or path.resolve().parent!=directory.resolve():raise ValueError('Ungültiger Ergebnispfad.')
  if path.stat().st_size>512*1024**2:raise ValueError('Spur zu groß.')
  return path.read_bytes()