"""Ephemeral admin chat tests using the same worker and leases as the API."""
import json
import math
import socket
import threading
import time
import uuid
from inference import InferenceError

class ChatTests:
 def __init__(self,profiles,worker,scheduler):
  self.profiles=profiles;self.worker=worker;self.scheduler=scheduler
  self.lock=threading.RLock();self.job=None;self.cancel=threading.Event();self.socket=None
 def status(self):
  with self.lock:job=dict(self.job) if self.job else None
  return dict(job=job,worker=self.worker.status(),scheduler=self.scheduler.status())
 def start(self,data):
  if set(data)!={'profile_id','messages','max_tokens'}:raise ValueError('Profil, Nachrichten und Antwortlimit erforderlich.')
  messages=data['messages'];limit=data['max_tokens']
  if type(limit)!=int or not 1<=limit<=4096:raise ValueError('Antwortlimit: 1–4096 Token.')
  if not isinstance(messages,list) or not 1<=len(messages)<=32:raise ValueError('1–32 Textnachrichten erforderlich.')
  for m in messages:
   if not isinstance(m,dict) or set(m)!={'role','content'} or m['role'] not in ('user','assistant','system') or not isinstance(m['content'],str) or not m['content'].strip():raise ValueError('Nur Textnachrichten mit Rolle user, assistant oder system unterstützt.')
  if sum(len(m['content']) for m in messages)>24000:raise ValueError('Testverlauf zu lang (maximal 24.000 Zeichen). Neuen Chat beginnen.')
  profile=next((p for p in self.profiles.status()['profiles'] if p['id']==data['profile_id'] and p['kind']=='chat'),None)
  if not profile or not profile['runnable']:raise ValueError('Kein ausführbares Chatprofil. Modell und CUDA-Build prüfen.')
  with self.lock:
   if self.job and self.job['state']=='running':raise ValueError('Ein Chat-Test läuft bereits.')
   self.cancel.clear();self.job=dict(id=uuid.uuid4().hex,state='running',phase='Wartet auf freie Modellreservierung',profile_id=profile['id'],profile_name=profile['name'],started_at=time.time(),answer='',reasoning='',error=None,timings={},usage={})
   threading.Thread(target=self._run,args=(profile,messages,limit),daemon=True).start()
   return dict(self.job)
 def phase(self,text):
  with self.lock:self.job['phase']=text
 def stop(self):
  self.cancel.set()
  with self.lock:sock=self.socket
  if sock:
   try:sock.shutdown(socket.SHUT_RDWR)
   except OSError:pass
  return {'cancellation_requested':True}
 def unload(self):
  if not self.scheduler.unload_idle():raise ValueError('Es laufen noch Anfragen. Erst deren Ende abwarten; andere Antworten werden nicht abgebrochen.')
  return self.status()
 def _run(self,profile,messages,limit):
  conn=None;outcome='failed';message='Chat-Test fehlgeschlagen.';error=None
  try:
   key=('chat',profile['id'],profile['revision'])
   def allowed():return not self.cancel.is_set() and any(p['id']==profile['id'] and p['revision']==profile['revision'] for p in self.profiles.status()['profiles'])
   def prepare():
    self.phase('Speicher wird geprüft · Profil wird geladen')
    self.worker.ensure(profile,cancel=self.cancel.is_set)
   with self.scheduler.lease(key,profile['parameters']['slots'],prepare,allowed=allowed):
    if self.cancel.is_set():raise InterruptedError()
    self.phase('Antwort wird erzeugt')
    body=dict(model=profile['name'],messages=messages,max_tokens=limit,stream=True,stream_options={"include_usage":True},**{k:profile['parameters'][k] for k in ('temperature','top_p','top_k')})
    conn,token=self.worker.connect()
    conn.request('POST','/v1/chat/completions',json.dumps(body).encode(),headers={'Content-Type':'application/json','Authorization':'Bearer '+token})
    with self.lock:self.socket=conn.sock
    if self.cancel.is_set():raise InterruptedError()
    response=conn.getresponse()
    if response.status!=200:
     raise InferenceError('llama.cpp hat den Test abgelehnt. Kontextbudget, Nachrichtenlänge und Profilparameter prüfen.')
    buffer=b'';total=0;done=False;deadline=time.monotonic()+600
    while time.monotonic()<deadline:
     if self.cancel.is_set():raise InterruptedError()
     chunk=response.read1(8192)
     if not chunk:break
     buffer+=chunk;total+=len(chunk)
     if total>4*1024**2:raise InferenceError('Testantwort überschreitet 4 MiB.')
     while b'\n' in buffer:
      line,buffer=buffer.split(b'\n',1);line=line.strip()
      if not line.startswith(b'data:'):continue
      payload=line[5:].strip()
      if payload==b'[DONE]':done=True;break
      event=json.loads(payload)
      if 'error' in event:raise InferenceError('Der Modellworker meldet einen Fehler während der Antwort.')
      with self.lock:
       for section,fields in [('timings',('prompt_n','prompt_ms','prompt_per_second','predicted_n','predicted_ms','predicted_per_second')),('usage',('prompt_tokens','completion_tokens','total_tokens'))]:
        values=event.get(section)
        if isinstance(values,dict):
         self.job[section]={k:v for k,v in values.items() if k in fields and type(v) in (int,float) and math.isfinite(v) and v>=0}
      for choice in event.get('choices',[]):
       delta=choice.get('delta',{})
       with self.lock:
        for source,target in [('content','answer'),('reasoning_content','reasoning')]:
         if isinstance(delta.get(source),str):self.job[target]+=delta[source]
        if choice.get('finish_reason'):self.job['finish_reason']=choice['finish_reason']
     if done:break
    if not done:raise InferenceError('Antwort unterbrochen oder Zeitlimit erreicht. '+(self.worker.status().get('error') or 'Ein GPU-OOM ist ohne eindeutigen Nachweis nicht bestätigt.'))
   outcome='complete';message='Antwort fertig · Modell bleibt für weitere Anfragen geladen'
  except Exception as exc:
   cancelled=self.cancel.is_set()
   message='Chat-Test abgebrochen.' if cancelled else (str(exc) if isinstance(exc,InferenceError) else 'Verbindung zum Modell unterbrochen. '+(self.worker.status().get('error') or 'Speicherfehler, Modellabsturz oder Zeitlimit möglich; Ursache nicht eindeutig.'))
   outcome='cancelled' if cancelled else 'failed';error=None if cancelled else message
  finally:
   if conn:conn.close()
   with self.lock:self.socket=None;self.job.update(state=outcome,error=error,phase=message,finished_at=time.time())